Dispatch to use fp32 distance computation in NN Descent depending on data dimensions

c/include/cuvs/neighbors/nn_descent.h

@@ -15,6 +15,18 @@
 extern "C" {
 #endif
 
+/**
+ * @brief Dtype to use for distance computation
+ * - `NND_DIST_COMP_AUTO`: Automatically determine the best dtype for distance computation based on the dataset dimensions.
+ * - `NND_DIST_COMP_FP32`: Use fp32 distance computation for better precision at the cost of performance and memory usage.
+ * - `NND_DIST_COMP_FP16`: Use fp16 distance computation.
+ */
+typedef enum {
+  NND_DIST_COMP_AUTO = 0,
+  NND_DIST_COMP_FP32 = 1,
+  NND_DIST_COMP_FP16 = 2
+} cuvsNNDescentDistCompDtype;
+
 /**
  * @defgroup nn_descent_c_index_params The nn-descent algorithm parameters.
  * @{
@@ -34,6 +46,8 @@ extern "C" {
  * `max_iterations`: The number of iterations that nn-descent will refine
  * the graph for. More iterations produce a better quality graph at cost of performance
  * `termination_threshold`: The delta at which nn-descent will terminate its iterations
+ * `return_distances`: Boolean to decide whether to return distances array
+ * `dist_comp_dtype`: dtype to use for distance computation. Defaults to `NND_DIST_COMP_AUTO` which automatically determines the best dtype for distance computation based on the dataset dimensions. Use `NND_DIST_COMP_FP32` for better precision at the cost of performance and memory usage. This option is only valid when data type is fp32. Use `NND_DIST_COMP_FP16` for better performance and memory usage at the cost of precision.
  */
 struct cuvsNNDescentIndexParams {
   cuvsDistanceType metric;
@@ -43,6 +57,7 @@ struct cuvsNNDescentIndexParams {
   size_t max_iterations;
   float termination_threshold;
   bool return_distances;
+  cuvsNNDescentDistCompDtype dist_comp_dtype;
 };
 
 typedef struct cuvsNNDescentIndexParams* cuvsNNDescentIndexParams_t;

c/src/neighbors/nn_descent.cpp

@@ -43,6 +43,7 @@ void* _build(cuvsResources_t res,
   build_params.max_iterations            = params.max_iterations;
   build_params.termination_threshold     = params.termination_threshold;
   build_params.return_distances          = params.return_distances;
+  build_params.dist_comp_dtype           = static_cast<cuvs::neighbors::nn_descent::DIST_COMP_DTYPE>(static_cast<int>(params.dist_comp_dtype));
 
   using graph_type = raft::host_matrix_view<IdxT, int64_t, raft::row_major>;
   std::optional<graph_type> graph;
@@ -177,7 +178,8 @@ extern "C" cuvsError_t cuvsNNDescentIndexParamsCreate(cuvsNNDescentIndexParams_t
       .intermediate_graph_degree = cpp_params.intermediate_graph_degree,
       .max_iterations            = cpp_params.max_iterations,
       .termination_threshold     = cpp_params.termination_threshold,
-      .return_distances          = cpp_params.return_distances};
+      .return_distances          = cpp_params.return_distances,
+      .dist_comp_dtype           = static_cast<cuvsNNDescentDistCompDtype>(static_cast<int>(cpp_params.dist_comp_dtype))};
   });
 }
 

cpp/include/cuvs/neighbors/nn_descent.hpp

@@ -24,6 +24,16 @@ namespace cuvs::neighbors::nn_descent {
  * @{
  */
 
+/**
+ * @brief Dtype to use for distance computation
+ * - `AUTO`: Automatically determine the best dtype for distance computation based on the dataset
+ * dimensions.
+ * - `FP32`: Use fp32 distance computation for better precision at the cost of performance and
+ * memory usage.
+ * - `FP16`: Use fp16 distance computation.
+ */
+enum class DIST_COMP_DTYPE { AUTO = 0, FP32 = 1, FP16 = 2 };
+
 /**
  * @brief Parameters used to build an nn-descent index
  * - `graph_degree`: For an input dataset of dimensions (N, D),
@@ -37,13 +47,19 @@ namespace cuvs::neighbors::nn_descent {
  * the graph for. More iterations produce a better quality graph at cost of performance
  * - `termination_threshold`: The delta at which nn-descent will terminate its iterations
  * - `return_distances`: Boolean to decide whether to return distances array
+ * - `dist_comp_dtype`: dtype to use for distance computation. Defaults to `AUTO` which
+ * automatically determines the best dtype for distance computation based on the dataset dimensions.
+ * Use `FP32` for better precision at the cost of performance and memory usage. This option is only
+ * valid when data type is fp32. Use `FP16` for better performance and memory usage at the cost of
+ * precision.
  */
 struct index_params : cuvs::neighbors::index_params {
   size_t graph_degree              = 64;
   size_t intermediate_graph_degree = 128;
   size_t max_iterations            = 20;
   float termination_threshold      = 0.0001;
   bool return_distances            = true;
+  DIST_COMP_DTYPE dist_comp_dtype  = DIST_COMP_DTYPE::AUTO;
 
   /** @brief Construct NN descent parameters for a specific kNN graph degree
    *