diff --git a/cpp/src/distance/detail/distance_ops/l2_exp.cuh b/cpp/src/distance/detail/distance_ops/l2_exp.cuh index c0433c676a..d059836ed1 100644 --- a/cpp/src/distance/detail/distance_ops/l2_exp.cuh +++ b/cpp/src/distance/detail/distance_ops/l2_exp.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -28,6 +28,23 @@ __device__ constexpr AccT get_clamp_precision() } } +/** + * Whether an expanded L2 distance is round-off error of a self-distance and should be clamped to + * zero. Self-neighboring points have (aNorm == bNorm) == accVal, but the dot product (accVal) can + * have round-off errors, which will cause (aNorm == bNorm) ~ accVal instead. + * + * The round-off error scales with the norms, so the tolerance is relative to them for squared + * norms below one. Otherwise, distinct points with equal norms and a small magnitude (for example, + * a spectral embedding) would have their distances incorrectly clamped to zero. + */ +template +DI bool is_self_distance_round_off(AccT dist, AccT aNorm, AccT bNorm) +{ + const AccT norm_product = aNorm * bNorm; + const AccT norm_scale = norm_product < AccT(1) ? norm_product : AccT(1); + return (aNorm == bNorm) && (dist * dist < get_clamp_precision() * norm_scale); +} + // Epilogue operator for CUTLASS based kernel template struct l2_exp_cutlass_op { @@ -38,13 +55,7 @@ struct l2_exp_cutlass_op { inline __device__ AccT operator()(AccT aNorm, AccT bNorm, AccT accVal) const noexcept { AccT outVal = aNorm + bNorm - AccT(2.0) * accVal; - - /** - * Self-neighboring points should have (aNorm == bNorm) == accVal and the dot product (accVal) - * can sometimes have round-off errors, which will cause (aNorm == bNorm) ~ accVal instead. - */ - outVal = - outVal * AccT(!((outVal * outVal < get_clamp_precision()) * (aNorm == bNorm))); + outVal = outVal * AccT(!is_self_distance_round_off(outVal, aNorm, bNorm)); return sqrt ? raft::sqrt(outVal * static_cast(outVal > AccT(0))) : outVal; } @@ -112,15 +123,9 @@ struct l2_exp_distance_op { for (int j = 0; j < Policy::AccColsPerTh; ++j) { AccT accVal = acc[i][j]; AccT val = regxn[i] + regyn[j] - (AccT)2.0 * accVal; - - /** - * Self-neighboring points should have (aNorm == bNorm) == accVal and the dot product - * (accVal) can sometimes have round-off errors, which will cause (aNorm == bNorm) ~ accVal - * instead. - */ - acc[i][j] = val * static_cast((val > AccT(0))) * - static_cast( - !((val * val < get_clamp_precision()) * (regxn[i] == regyn[j]))); + acc[i][j] = + val * static_cast((val > AccT(0))) * + static_cast(!is_self_distance_round_off(val, regxn[i], regyn[j])); } } if (sqrt) { diff --git a/cpp/tests/distance/dist_l2_exp.cu b/cpp/tests/distance/dist_l2_exp.cu index a91daff1b3..1c9dbe4e28 100644 --- a/cpp/tests/distance/dist_l2_exp.cu +++ b/cpp/tests/distance/dist_l2_exp.cu @@ -1,11 +1,17 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "../test_utils.cuh" #include "distance_base.cuh" +#include +#include + +#include +#include + namespace cuvs { namespace distance { @@ -162,5 +168,46 @@ INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceEucExpTestD, ::testing::ValuesIn( class BigMatrixEucExp : public BigMatrixDistanceTest {}; TEST_F(BigMatrixEucExp, Result) {} + +// Distinct points with equal norms must keep their (small) distances. The clamp of self-distance +// round-off to zero used to apply an absolute tolerance, which collapsed such points when the data +// has a small magnitude (e.g. a spectral embedding). +template +void test_small_scale_equal_norms(T scale) +{ + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr int64_t n = 8; + std::vector h_x(n * n, T{0}); + for (int64_t i = 0; i < n; i++) { + h_x[i * n + i] = scale; + } + auto x = raft::make_device_matrix(handle, n, n); + auto d = raft::make_device_matrix(handle, n, n); + raft::update_device(x.data_handle(), h_x.data(), h_x.size(), stream); + + for (auto metric : {DistanceType::L2Expanded, DistanceType::L2SqrtExpanded}) { + cuvs::distance::pairwise_distance(handle, + raft::make_const_mdspan(x.view()), + raft::make_const_mdspan(x.view()), + d.view(), + metric); + std::vector h_d(n * n); + raft::update_host(h_d.data(), d.data_handle(), h_d.size(), stream); + raft::resource::sync_stream(handle, stream); + + const T expected = + metric == DistanceType::L2Expanded ? T(2) * scale * scale : std::sqrt(T(2)) * scale; + for (int64_t i = 0; i < n; i++) { + for (int64_t j = 0; j < n; j++) { + EXPECT_NEAR(h_d[i * n + j], i == j ? T(0) : expected, T(1e-3) * expected) + << "metric=" << static_cast(metric) << " i=" << i << " j=" << j; + } + } + } +} + +TEST(DistanceEucExpSmallScale, EqualNormsF) { test_small_scale_equal_norms(1e-2f); } +TEST(DistanceEucExpSmallScale, EqualNormsD) { test_small_scale_equal_norms(1e-5); } } // end namespace distance } // namespace cuvs diff --git a/cpp/tests/neighbors/distance_nn.cu b/cpp/tests/neighbors/distance_nn.cu index 5ac80fda73..9ab376927c 100644 --- a/cpp/tests/neighbors/distance_nn.cu +++ b/cpp/tests/neighbors/distance_nn.cu @@ -16,8 +16,10 @@ #include #include #include +#include #include +#include namespace cuvs::neighbors { @@ -446,6 +448,77 @@ TEST(Top1nnPlan, RejectsMismatchedLaunch) plan))); } +// Distinct points with equal norms and a small magnitude (as in a spectral embedding) must not be +// treated as self-neighbors when finding the nearest neighbor. +TEST(Top1nnEqualNorms, SmallScale) +{ + using Backend = cuvs::distance::detail::Top1nnBackend; + using Kvp = raft::KeyValuePair; + raft::resources handle; + auto stream = raft::resource::get_cuda_stream(handle); + constexpr int m = 1000; + constexpr int n = 5; + constexpr int k = 5; + constexpr float s = 1e-2f; + std::vector h_x(m * k, 0.0f); + std::vector h_y(n * k, 0.0f); + for (int i = 0; i < m; i++) { + h_x[i * k + i % n] = s; + } + for (int j = 0; j < n; j++) { + h_y[j * k + j] = s; + } + auto x = raft::make_device_matrix(handle, m, k); + auto y = raft::make_device_matrix(handle, n, k); + auto x_norm = raft::make_device_vector(handle, m); + auto y_norm = raft::make_device_vector(handle, n); + auto output = raft::make_device_vector(handle, m); + raft::update_device(x.data_handle(), h_x.data(), h_x.size(), stream.get()); + raft::update_device(y.data_handle(), h_y.data(), h_y.size(), stream.get()); + raft::linalg::rowNorm( + x_norm.data_handle(), x.data_handle(), k, m, stream.get()); + raft::linalg::rowNorm( + y_norm.data_handle(), y.data_handle(), k, n, stream.get()); + + cuvs::distance::detail::Top1nnTuning tuning{}; + for (auto backend : {Backend::Cutlass, Backend::Unfused}) { + if (!cuvs::distance::detail::is_top_1_nn_backend_available( + backend, x.data_handle(), y.data_handle(), m, n, k, DistanceType::L2Expanded)) { + continue; + } + const auto plan = cuvs::distance::probe_top_1_nn( + handle, x.data_handle(), y.data_handle(), m, n, k, tuning, DistanceType::L2Expanded, backend); + ASSERT_TRUE(plan.available); + auto workspace = raft::make_device_vector(handle, plan.workspace_bytes); + cuvs::distance::top_1_nn(handle, + output.data_handle(), + x.data_handle(), + y.data_handle(), + x_norm.data_handle(), + y_norm.data_handle(), + m, + n, + k, + tuning, + workspace.data_handle(), + workspace.size(), + false, + true, + true, + DistanceType::L2Expanded, + 0.0f, + plan); + std::vector h_out(m); + raft::update_host(h_out.data(), output.data_handle(), m, stream.get()); + raft::resource::sync_stream(handle, stream); + int n_wrong = 0; + for (int i = 0; i < m; i++) { + n_wrong += h_out[i].key != i % n; + } + EXPECT_EQ(n_wrong, 0) << "backend=" << static_cast(backend); + } +} + TEST(Top1nnPlan, ExplicitUnavailableBackendIsStrict) { raft::resources handle;