Fix L2 expanded distance clamp for cutlass and legacy fused backend

cpp/src/distance/detail/distance.cuh

@@ -7,21 +7,35 @@
 
 #include "distance_ops/all_ops.cuh"
 #include "pairwise_matrix/dispatch.cuh"
+
 #include <cuvs/distance/distance.hpp>
+
 #include <raft/core/operators.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/linalg/gemm.cuh>
 #include <raft/linalg/map.cuh>
 #include <raft/linalg/norm.cuh>
 #include <raft/linalg/reduce.cuh>
 #include <raft/util/cuda_dev_essentials.cuh>  // to_float
+#include <raft/util/kernel_launch.hpp>
 
+#include <cstdint>
 #include <type_traits>
 
 namespace cuvs {
 namespace distance {
 namespace detail {
 
+template <typename OutT, typename AccT, typename FinOpT, typename IdxT>
+RAFT_KERNEL set_self_distances_kernel(OutT* out,
+                                      IdxT diagonal_size,
+                                      IdxT leading_dimension,
+                                      FinOpT fin_op)
+{
+  IdxT idx = static_cast<IdxT>(blockIdx.x) * blockDim.x + threadIdx.x;
+  if (idx < diagonal_size) { out[std::size_t(idx) * leading_dimension + idx] = fin_op(AccT{0}, 0); }
+}
+
 /**
  * @brief: A tag type for overload resolution based on DistanceType
  *
@@ -591,6 +605,23 @@ void distance_impl_l2_expanded(  // NOTE: different name
   ops::l2_exp_distance_op<DataT, AccT, IdxT> distance_op{perform_sqrt};
   pairwise_matrix_dispatch<decltype(distance_op), DataT, AccT, OutT, FinOpT, IdxT>(
     distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major);
+
+  // When both inputs are the same matrix, diagonal entries compare each row with itself.
+  // Set those known self-distances to zero.
+  if (x == y && m > 0 && n > 0 && (is_row_major || m == n)) {
+    constexpr uint32_t block_size = 256;
+    IdxT diagonal_size            = std::min(m, n);
+    IdxT leading_dimension        = is_row_major ? n : m;
+    auto grid_size                = raft::ceildiv(diagonal_size, IdxT{block_size});
+    raft::launch_kernel(stream,
+                        grid_size,
+                        block_size,
+                        set_self_distances_kernel<OutT, AccT, FinOpT, IdxT>,
+                        out,
+                        diagonal_size,
+                        leading_dimension,
+                        fin_op);
+  }
 }
 
 template <typename DataT, typename AccT, typename OutT, typename FinOpT, typename IdxT = int>

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
  */
 
@@ -12,22 +12,6 @@
 
 namespace cuvs::distance::detail::ops {
 
-/**
- * Reserve 1 digit of precision from each floating-point type
- * for round-off error tolerance.
- * @tparam DataT
- */
-template <typename DataT, typename AccT>
-__device__ constexpr AccT get_clamp_precision()
-{
-  switch (sizeof(DataT)) {
-    case 2: return AccT{1e-3};
-    case 4: return AccT{1e-6};
-    case 8: return AccT{1e-15};
-    default: return AccT{0};
-  }
-}
-
 // Epilogue operator for CUTLASS based kernel
 template <typename DataT, typename AccT>
 struct l2_exp_cutlass_op {
@@ -39,13 +23,9 @@ struct l2_exp_cutlass_op {
   {
     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<DataT, AccT>()) * (aNorm == bNorm)));
-    return sqrt ? raft::sqrt(outVal * static_cast<AccT>(outVal > AccT(0))) : outVal;
+    // Expanded L2 can be slightly negative due to round-off; clamp to zero.
+    outVal = outVal < AccT{0} ? AccT{0} : outVal;
+    return sqrt ? raft::sqrt(outVal) : outVal;
   }
 
   __device__ AccT operator()(DataT aData) const noexcept
@@ -113,14 +93,8 @@ struct l2_exp_distance_op {
         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<AccT>((val > AccT(0))) *
-                    static_cast<AccT>(
-                      !((val * val < get_clamp_precision<DataT, AccT>()) * (regxn[i] == regyn[j])));
+        // Expanded L2 can be slightly negative due to round-off; clamp to zero.
+        acc[i][j] = val < AccT{0} ? AccT{0} : val;
       }
     }
     if (sqrt) {

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 <raft/core/device_mdarray.hpp>
+#include <raft/util/cudart_utils.hpp>
+
+#include <cmath>
+#include <vector>
+
 namespace cuvs {
 namespace distance {
 
@@ -162,5 +168,92 @@ INSTANTIATE_TEST_CASE_P(DistanceTests, DistanceEucExpTestD, ::testing::ValuesIn(
 
 class BigMatrixEucExp : public BigMatrixDistanceTest<cuvs::distance::DistanceType::L2Expanded> {};
 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 <typename T>
+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<T> 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<T, int64_t>(handle, n, n);
+  auto d = raft::make_device_matrix<T, int64_t>(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<T> 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<int>(metric) << " i=" << i << " j=" << j;
+      }
+    }
+  }
+}
+
+TEST(DistanceEucExpSmallScale, EqualNormsF) { test_small_scale_equal_norms<float>(1e-2f); }
+TEST(DistanceEucExpSmallScale, EqualNormsD) { test_small_scale_equal_norms<double>(1e-5); }
+
+template <typename Layout>
+void test_exact_self_distance_diagonal()
+{
+  raft::resources handle;
+  auto stream              = raft::resource::get_cuda_stream(handle);
+  constexpr int64_t n_rows = 50;
+  constexpr int64_t n_cols = 10;
+  constexpr bool row_major = std::is_same_v<Layout, raft::row_major>;
+  std::vector<float> h_x(n_rows * n_cols);
+  for (int64_t i = 0; i < n_rows; ++i) {
+    for (int64_t j = 0; j < n_cols; ++j) {
+      auto index = row_major ? i * n_cols + j : j * n_rows + i;
+      h_x[index] = static_cast<float>(((i + 3) * (j + 5)) % 37) / 37.0f;
+    }
+  }
+
+  auto x = raft::make_device_matrix<float, int64_t, Layout>(handle, n_rows, n_cols);
+  auto d = raft::make_device_matrix<float, int64_t, Layout>(handle, n_rows, n_rows);
+  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<float> h_d(n_rows * n_rows);
+    raft::update_host(h_d.data(), d.data_handle(), h_d.size(), stream);
+    raft::resource::sync_stream(handle, stream);
+
+    for (int64_t i = 0; i < n_rows; ++i) {
+      EXPECT_EQ(h_d[i * n_rows + i], 0.0f) << "metric=" << static_cast<int>(metric) << " row=" << i;
+    }
+  }
+}
+
+TEST(DistanceEucExpSelfDistance, ExactDiagonalRowMajor)
+{
+  test_exact_self_distance_diagonal<raft::row_major>();
+}
+
+TEST(DistanceEucExpSelfDistance, ExactDiagonalColMajor)
+{
+  test_exact_self_distance_diagonal<raft::col_major>();
+}
 }  // end namespace distance
 }  // namespace cuvs

cpp/tests/neighbors/brute_force.cu

@@ -18,6 +18,7 @@
 #include <cuda_fp16.h>
 
 #include <algorithm>
+#include <cmath>
 #include <vector>
 
 namespace cuvs::neighbors::brute_force {
@@ -195,6 +196,83 @@ TEST_P(KNNTest_half_int64_t, BruteForce) { this->testBruteForce(); }
 INSTANTIATE_TEST_CASE_P(KNNTest, KNNTest_float_int64_t, ::testing::ValuesIn(inputs<float>));
 INSTANTIATE_TEST_CASE_P(KNNTest, KNNTest_half_int64_t, ::testing::ValuesIn(inputs<half>));
 
+TEST(BruteForceEqualNorms, PreservesNeighborOrdering)
+{
+  raft::resources handle;
+  auto stream                  = raft::resource::get_cuda_stream(handle);
+  constexpr int64_t n_database = 128;
+  constexpr int64_t n_queries  = 1;
+  constexpr int64_t dim        = 2;
+  constexpr double angle_step  = 2e-3;
+
+  std::vector<float> h_database(n_database * dim);
+  for (int64_t i = 0; i < n_database; ++i) {
+    auto rank               = (i * 73) % n_database;
+    double angle            = static_cast<double>(rank) * angle_step;
+    h_database[i * dim]     = static_cast<float>(std::cos(angle));
+    h_database[i * dim + 1] = static_cast<float>(std::sin(angle));
+  }
+  std::vector<float> h_queries{1.0f, 0.0f};
+
+  struct reference_neighbor {
+    double squared_distance;
+    int64_t index;
+  };
+  std::vector<reference_neighbor> reference(n_database);
+  for (int64_t i = 0; i < n_database; ++i) {
+    double dx    = static_cast<double>(h_queries[0]) - static_cast<double>(h_database[i * dim]);
+    double dy    = static_cast<double>(h_queries[1]) - static_cast<double>(h_database[i * dim + 1]);
+    reference[i] = {.squared_distance = dx * dx + dy * dy, .index = i};
+  }
+  std::ranges::sort(
+    reference, [](const reference_neighbor& lhs, const reference_neighbor& rhs) -> bool {
+      return lhs.squared_distance < rhs.squared_distance ||
+             (lhs.squared_distance == rhs.squared_distance && lhs.index < rhs.index);
+    });
+
+  auto database = raft::make_device_matrix<float, int64_t>(handle, n_database, dim);
+  auto queries  = raft::make_device_matrix<float, int64_t>(handle, n_queries, dim);
+  raft::update_device(database.data_handle(), h_database.data(), h_database.size(), stream);
+  raft::update_device(queries.data_handle(), h_queries.data(), h_queries.size(), stream);
+
+  for (auto metric :
+       {cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::L2SqrtExpanded}) {
+    index_params index_params;
+    index_params.metric = metric;
+    auto index          = build(handle, index_params, raft::make_const_mdspan(database.view()));
+
+    for (int64_t k : {int64_t{15}, int64_t{65}}) {
+      auto neighbors = raft::make_device_matrix<int64_t, int64_t>(handle, n_queries, k);
+      auto distances = raft::make_device_matrix<float, int64_t>(handle, n_queries, k);
+      search(handle,
+             search_params{},
+             index,
+             raft::make_const_mdspan(queries.view()),
+             neighbors.view(),
+             distances.view(),
+             filtering::none_sample_filter{});
+
+      std::vector<int64_t> h_neighbors(k);
+      std::vector<float> h_distances(k);
+      raft::update_host(h_neighbors.data(), neighbors.data_handle(), k, stream);
+      raft::update_host(h_distances.data(), distances.data_handle(), k, stream);
+      raft::resource::sync_stream(handle, stream);
+
+      for (int64_t i = 0; i < k; ++i) {
+        auto expected_distance = metric == cuvs::distance::DistanceType::L2Expanded
+                                   ? reference[i].squared_distance
+                                   : std::sqrt(reference[i].squared_distance);
+        EXPECT_EQ(h_neighbors[i], reference[i].index)
+          << "metric=" << static_cast<int>(metric) << " k=" << k << " rank=" << i;
+        EXPECT_NEAR(h_distances[i],
+                    expected_distance,
+                    metric == cuvs::distance::DistanceType::L2Expanded ? 5e-7 : 5e-5)
+          << "metric=" << static_cast<int>(metric) << " k=" << k << " rank=" << i;
+      }
+    }
+  }
+}
+
 // Also test with larger random inputs, including col-major inputs
 struct RandomKNNInputs {
   int num_queries;