Moving dice distance from raft to cuvs and support for half types

cpp/CMakeLists.txt

@@ -379,6 +379,9 @@ if(NOT BUILD_CPU_ONLY)
     src/distance/detail/pairwise_matrix/dispatch_cosine_float_float_float_int.cu
     src/distance/detail/pairwise_matrix/dispatch_cosine_half_float_float_int.cu
     src/distance/detail/pairwise_matrix/dispatch_cosine_double_double_double_int.cu
+    src/distance/detail/pairwise_matrix/dispatch_dice_float_float_float_int.cu
+    src/distance/detail/pairwise_matrix/dispatch_dice_half_float_float_int.cu
+    src/distance/detail/pairwise_matrix/dispatch_dice_double_double_double_int.cu
     src/distance/detail/pairwise_matrix/dispatch_hamming_unexpanded_float_float_float_int.cu
     src/distance/detail/pairwise_matrix/dispatch_hamming_unexpanded_half_float_float_int.cu
     src/distance/detail/pairwise_matrix/dispatch_hamming_unexpanded_double_double_double_int.cu

cpp/src/distance/detail/distance.cuh

@@ -45,6 +45,7 @@ using distance_tag = std::integral_constant<DistanceType, d>;
  * - DistanceType::Canberra:
  * - DistanceType::CorrelationExpanded:
  * - DistanceType::CosineExpanded:
+ * - DistanceType::DiceExpanded:
  * - DistanceType::HammingUnexpanded:
  * - DistanceType::HellingerExpanded:
  * - DistanceType::JensenShannon:
@@ -215,6 +216,59 @@ void distance_impl(raft::resources const& handle,
     distance_op, m, n, k, x, y, x_norm, y_norm, out, fin_op, stream, is_row_major);
 }
 
+template <typename DataT, typename AccT, typename OutT, typename FinOpT, typename IdxT = int>
+void distance_impl(raft::resources const& handle,
+                   distance_tag<DistanceType::DiceExpanded> distance_type,
+                   const DataT* x,
+                   const DataT* y,
+                   OutT* out,
+                   IdxT m,
+                   IdxT n,
+                   IdxT k,
+                   AccT* workspace,
+                   size_t worksize,
+                   FinOpT fin_op,
+                   bool is_row_major,
+                   DataT)  // unused
+{
+  // raft distance support inputs as float/double and output as uint8_t/float/double.
+  static_assert(!((sizeof(OutT) > 1) && (sizeof(AccT) != sizeof(OutT))),
+                "OutT can be uint8_t, float, double,"
+                "if sizeof(OutT) > 1 then sizeof(AccT) == sizeof(OutT).");
+
+  ASSERT(!(worksize < (m + n) * sizeof(AccT)), "workspace size error");
+  ASSERT(workspace != nullptr, "workspace is null");
+
+  cudaStream_t stream = raft::resource::get_cuda_stream(handle);
+
+  AccT* x_norm = workspace;
+  AccT* y_norm = workspace;
+  // TODO: Column major case looks to have lower accuracy for X == Y,
+  // perhaps the use of stridedSummationKernel could be causing this,
+  // need to investigate and fix.
+  if (x == y && is_row_major) {
+    raft::linalg::reduce<true, true>(
+      x_norm, x, k, std::max(m, n), (AccT)0, stream, false, raft::identity_op(), raft::add_op());
+  } else {
+    y_norm += m;
+    if (is_row_major) {
+      raft::linalg::reduce<true, true>(
+        x_norm, x, k, m, (AccT)0, stream, false, raft::identity_op(), raft::add_op());
+      raft::linalg::reduce<true, true>(
+        y_norm, y, k, n, (AccT)0, stream, false, raft::identity_op(), raft::add_op());
+    } else {
+      raft::linalg::reduce<false, true>(
+        x_norm, x, k, m, (AccT)0, stream, false, raft::identity_op(), raft::add_op());
+      raft::linalg::reduce<false, true>(
+        y_norm, y, k, n, (AccT)0, stream, false, raft::identity_op(), raft::add_op());
+    }
+  }
+
+  ops::dice_distance_op<DataT, AccT, IdxT> distance_op{};
+  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);
+}
+
 template <typename DataT, typename AccT, typename OutT, typename FinOpT, typename IdxT = int>
 void distance_impl(raft::resources const& handle,
                    distance_tag<DistanceType::HammingUnexpanded> distance_type,
@@ -812,9 +866,11 @@ template <cuvs::distance::DistanceType distanceType,
           typename Index_ = int>
 size_t getWorkspaceSize(const InType* x, const InType* y, Index_ m, Index_ n, Index_ k)
 {
-  size_t worksize             = 0;
-  constexpr bool is_allocated = (distanceType <= cuvs::distance::DistanceType::CosineExpanded) ||
-                                (distanceType == cuvs::distance::DistanceType::CorrelationExpanded);
+  size_t worksize = 0;
+  constexpr bool is_allocated =
+    (distanceType <= cuvs::distance::DistanceType::CosineExpanded) ||
+    (distanceType == cuvs::distance::DistanceType::CorrelationExpanded) ||
+    (distanceType == cuvs::distance::DistanceType::DiceExpanded);
   constexpr int numOfBuffers =
     (distanceType == cuvs::distance::DistanceType::CorrelationExpanded) ? 2 : 1;
 

cpp/src/distance/detail/distance_ops/all_ops.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -12,6 +12,7 @@
 #include "../distance_ops/canberra.cuh"
 #include "../distance_ops/correlation.cuh"
 #include "../distance_ops/cosine.cuh"
+#include "../distance_ops/dice.cuh"
 #include "../distance_ops/hamming.cuh"
 #include "../distance_ops/hellinger.cuh"
 #include "../distance_ops/jensen_shannon.cuh"

cpp/src/distance/detail/distance_ops/dice.cuh

@@ -0,0 +1,65 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <raft/util/cuda_dev_essentials.cuh>  // DI
+
+namespace cuvs::distance::detail::ops {
+
+/**
+ * @brief the expanded dice distance matrix calculation
+ *
+ * It computes the following equation:
+ *
+ * d(x, y) = 1 - 2*(x ⋅ y) / ( Σ(x) + Σ(y) )
+ */
+template <typename DataType, typename AccType, typename IdxType>
+struct dice_distance_op {
+  using DataT = DataType;
+  using AccT  = AccType;
+  using IdxT  = IdxType;
+
+  // Load norms of input data
+  static constexpr bool use_norms = true;
+  // Whether the core function requires so many instructions that it makes sense
+  // to reduce loop unrolling, etc. We do this to keep compile times in check.
+  static constexpr bool expensive_inner_loop = false;
+
+  // Size of shared memory. This is normally decided by the kernel policy, but
+  // some ops such as correlation_distance_op use more.
+  template <typename Policy>
+  static constexpr size_t shared_mem_size()
+  {
+    return Policy::SmemSize + ((Policy::Mblk + Policy::Nblk) * sizeof(AccT));
+  }
+
+  DI void core(AccT& acc, DataT& x, DataT& y) const
+  {
+    if constexpr ((std::is_same_v<AccT, float> && std::is_same_v<DataT, half>)) {
+      acc += __half2float(x) * __half2float(y);
+    } else {
+      acc += x * y;
+    }
+  };
+
+  template <typename Policy>
+  DI void epilog(AccT acc[Policy::AccRowsPerTh][Policy::AccColsPerTh],
+                 AccT* regxn,
+                 AccT* regyn,
+                 IdxT gridStrideX,
+                 IdxT gridStrideY) const
+  {
+#pragma unroll
+    for (int i = 0; i < Policy::AccRowsPerTh; ++i) {
+#pragma unroll
+      for (int j = 0; j < Policy::AccColsPerTh; ++j) {
+        acc[i][j] = 1.0 - (2 * acc[i][j] / (regxn[i] + regyn[j]));
+      }
+    }
+  }
+};
+
+}  // namespace cuvs::distance::detail::ops

cpp/src/distance/detail/pairwise_matrix/dispatch-ext.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -84,6 +84,8 @@ instantiate_cuvs_distance_detail_pairwise_matrix_dispatch_by_algo_default(
   cuvs::distance::detail::ops::correlation_distance_op, int);
 instantiate_cuvs_distance_detail_pairwise_matrix_dispatch_by_algo_default(
   cuvs::distance::detail::ops::cosine_distance_op, int);
+instantiate_cuvs_distance_detail_pairwise_matrix_dispatch_by_algo_default(
+  cuvs::distance::detail::ops::dice_distance_op, int);
 instantiate_cuvs_distance_detail_pairwise_matrix_dispatch_by_algo_default(
   cuvs::distance::detail::ops::hamming_distance_op, int);
 instantiate_cuvs_distance_detail_pairwise_matrix_dispatch_by_algo_default(

cpp/src/distance/detail/pairwise_matrix/dispatch_00_generate.py

@@ -81,6 +81,11 @@
         OpT="cuvs::distance::detail::ops::cosine_distance_op",
         archs=[60, 80],
     ),
+    dict(
+        path_prefix="dice",
+        OpT="cuvs::distance::detail::ops::dice_distance_op",
+        archs=[60, 80],
+    ),
     dict(
         path_prefix="hamming_unexpanded",
         OpT="cuvs::distance::detail::ops::hamming_distance_op",

cpp/src/distance/detail/pairwise_matrix/dispatch_dice_double_double_double_int.cu

@@ -0,0 +1,40 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+/*
+ * NOTE: this file is generated by dispatch_00_generate.py
+ *
+ * Make changes there and run in this directory:
+ *
+ * > python dispatch_00_generate.py
+ *
+ */
+
+#include "../distance_ops/all_ops.cuh"  // ops::*
+#include "dispatch-inl.cuh"             // dispatch
+#include "dispatch_sm60.cuh"
+#include "dispatch_sm80.cuh"
+#include <raft/core/operators.hpp>  // raft::identity_op
+#define instantiate_raft_distance_detail_pairwise_matrix_dispatch(                     \
+  OpT, DataT, AccT, OutT, FinOpT, IdxT)                                                \
+  template void cuvs::distance::detail::                                               \
+    pairwise_matrix_dispatch<OpT<DataT, AccT, IdxT>, DataT, AccT, OutT, FinOpT, IdxT>( \
+      OpT<DataT, AccT, IdxT> distance_op,                                              \
+      IdxT m,                                                                          \
+      IdxT n,                                                                          \
+      IdxT k,                                                                          \
+      const DataT* x,                                                                  \
+      const DataT* y,                                                                  \
+      const OutT* x_norm,                                                              \
+      const OutT* y_norm,                                                              \
+      OutT* out,                                                                       \
+      FinOpT fin_op,                                                                   \
+      cudaStream_t stream,                                                             \
+      bool is_row_major)
+
+instantiate_raft_distance_detail_pairwise_matrix_dispatch(
+  cuvs::distance::detail::ops::dice_distance_op, double, double, double, raft::identity_op, int);
+
+#undef instantiate_raft_distance_detail_pairwise_matrix_dispatch

cpp/src/distance/detail/pairwise_matrix/dispatch_dice_float_float_float_int.cu

@@ -0,0 +1,40 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+/*
+ * NOTE: this file is generated by dispatch_00_generate.py
+ *
+ * Make changes there and run in this directory:
+ *
+ * > python dispatch_00_generate.py
+ *
+ */
+
+#include "../distance_ops/all_ops.cuh"  // ops::*
+#include "dispatch-inl.cuh"             // dispatch
+#include "dispatch_sm60.cuh"
+#include "dispatch_sm80.cuh"
+#include <raft/core/operators.hpp>  // raft::identity_op
+#define instantiate_raft_distance_detail_pairwise_matrix_dispatch(                     \
+  OpT, DataT, AccT, OutT, FinOpT, IdxT)                                                \
+  template void cuvs::distance::detail::                                               \
+    pairwise_matrix_dispatch<OpT<DataT, AccT, IdxT>, DataT, AccT, OutT, FinOpT, IdxT>( \
+      OpT<DataT, AccT, IdxT> distance_op,                                              \
+      IdxT m,                                                                          \
+      IdxT n,                                                                          \
+      IdxT k,                                                                          \
+      const DataT* x,                                                                  \
+      const DataT* y,                                                                  \
+      const OutT* x_norm,                                                              \
+      const OutT* y_norm,                                                              \
+      OutT* out,                                                                       \
+      FinOpT fin_op,                                                                   \
+      cudaStream_t stream,                                                             \
+      bool is_row_major)
+
+instantiate_raft_distance_detail_pairwise_matrix_dispatch(
+  cuvs::distance::detail::ops::dice_distance_op, float, float, float, raft::identity_op, int);
+
+#undef instantiate_raft_distance_detail_pairwise_matrix_dispatch

cpp/src/distance/detail/pairwise_matrix/dispatch_dice_half_float_float_int.cu

@@ -0,0 +1,40 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+/*
+ * NOTE: this file is generated by dispatch_00_generate.py
+ *
+ * Make changes there and run in this directory:
+ *
+ * > python dispatch_00_generate.py
+ *
+ */
+
+#include "../distance_ops/all_ops.cuh"  // ops::*
+#include "dispatch-inl.cuh"             // dispatch
+#include "dispatch_sm60.cuh"
+#include "dispatch_sm80.cuh"
+#include <raft/core/operators.hpp>  // raft::identity_op
+#define instantiate_raft_distance_detail_pairwise_matrix_dispatch(                     \
+  OpT, DataT, AccT, OutT, FinOpT, IdxT)                                                \
+  template void cuvs::distance::detail::                                               \
+    pairwise_matrix_dispatch<OpT<DataT, AccT, IdxT>, DataT, AccT, OutT, FinOpT, IdxT>( \
+      OpT<DataT, AccT, IdxT> distance_op,                                              \
+      IdxT m,                                                                          \
+      IdxT n,                                                                          \
+      IdxT k,                                                                          \
+      const DataT* x,                                                                  \
+      const DataT* y,                                                                  \
+      const OutT* x_norm,                                                              \
+      const OutT* y_norm,                                                              \
+      OutT* out,                                                                       \
+      FinOpT fin_op,                                                                   \
+      cudaStream_t stream,                                                             \
+      bool is_row_major)
+
+instantiate_raft_distance_detail_pairwise_matrix_dispatch(
+  cuvs::distance::detail::ops::dice_distance_op, half, float, float, raft::identity_op, int);
+
+#undef instantiate_raft_distance_detail_pairwise_matrix_dispatch

cpp/src/distance/distance-ext.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2018-2025, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -212,6 +212,7 @@ void pairwise_distance(raft::resources const& handle,
 instantiate_cuvs_distance_distance_by_algo(cuvs::distance::DistanceType::Canberra);
 instantiate_cuvs_distance_distance_by_algo(cuvs::distance::DistanceType::CorrelationExpanded);
 instantiate_cuvs_distance_distance_by_algo(cuvs::distance::DistanceType::CosineExpanded);
+instantiate_cuvs_distance_distance_by_algo(cuvs::distance::DistanceType::DiceExpanded);
 instantiate_cuvs_distance_distance_by_algo(cuvs::distance::DistanceType::HammingUnexpanded);
 
 instantiate_cuvs_distance_distance_by_algo(cuvs::distance::DistanceType::HellingerExpanded);
@@ -299,6 +300,7 @@ instantiate_cuvs_distance_getWorkspaceSize_by_algo(cuvs::distance::DistanceType:
 instantiate_cuvs_distance_getWorkspaceSize_by_algo(
   cuvs::distance::DistanceType::CorrelationExpanded);
 instantiate_cuvs_distance_getWorkspaceSize_by_algo(cuvs::distance::DistanceType::CosineExpanded);
+instantiate_cuvs_distance_getWorkspaceSize_by_algo(cuvs::distance::DistanceType::DiceExpanded);
 instantiate_cuvs_distance_getWorkspaceSize_by_algo(cuvs::distance::DistanceType::HammingUnexpanded);
 
 instantiate_cuvs_distance_getWorkspaceSize_by_algo(cuvs::distance::DistanceType::HellingerExpanded);

cpp/src/distance/distance-inl.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2018-2025, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -296,6 +296,9 @@ void pairwise_distance(raft::resources const& handle,
     case DistanceType::RusselRaoExpanded:
       dispatch(std::integral_constant<DistanceType, DistanceType::RusselRaoExpanded>{});
       break;
+    case DistanceType::DiceExpanded:
+      dispatch(std::integral_constant<DistanceType, DistanceType::DiceExpanded>{});
+      break;
     default: THROW("Unknown or unsupported distance metric '%d'!", (int)metric);
   };
 }

cpp/src/distance/distance.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2018-2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2018-2025, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -77,6 +77,7 @@
 instantiate_cuvs_distance_distance_by_algo(cuvs::distance::DistanceType::Canberra);
 instantiate_cuvs_distance_distance_by_algo(cuvs::distance::DistanceType::CorrelationExpanded);
 instantiate_cuvs_distance_distance_by_algo(cuvs::distance::DistanceType::CosineExpanded);
+instantiate_cuvs_distance_distance_by_algo(cuvs::distance::DistanceType::DiceExpanded);
 instantiate_cuvs_distance_distance_by_algo(cuvs::distance::DistanceType::HammingUnexpanded);
 
 instantiate_cuvs_distance_distance_by_algo(cuvs::distance::DistanceType::HellingerExpanded);
@@ -165,6 +166,7 @@ instantiate_cuvs_distance_getWorkspaceSize_by_algo(cuvs::distance::DistanceType:
 instantiate_cuvs_distance_getWorkspaceSize_by_algo(
   cuvs::distance::DistanceType::CorrelationExpanded);
 instantiate_cuvs_distance_getWorkspaceSize_by_algo(cuvs::distance::DistanceType::CosineExpanded);
+instantiate_cuvs_distance_getWorkspaceSize_by_algo(cuvs::distance::DistanceType::DiceExpanded);
 instantiate_cuvs_distance_getWorkspaceSize_by_algo(cuvs::distance::DistanceType::HammingUnexpanded);
 
 instantiate_cuvs_distance_getWorkspaceSize_by_algo(cuvs::distance::DistanceType::HellingerExpanded);

cpp/tests/CMakeLists.txt

@@ -295,6 +295,7 @@ ConfigureTest(
   PATH distance/dist_canberra.cu
        distance/dist_correlation.cu
        distance/dist_cos.cu
+       distance/dist_dice.cu
        distance/dist_hamming.cu
        distance/dist_hellinger.cu
        distance/dist_inner_product.cu