fix type template and bring back raft_expects

rapidsai · rapids-bot · Jun 14, 2024 · May 29, 2024 · May 29, 2024 · May 29, 2024
commit 84abd964aa2c3dea25bcb64b6dee36618cbe684a
diff --git a/cpp/include/raft/neighbors/detail/cagra/cagra_build.cuh b/cpp/include/raft/neighbors/detail/cagra/cagra_build.cuh
@@ -240,7 +240,7 @@ void build_knn_graph(raft::resources const& res,
                      raft::host_matrix_view<IdxT, int64_t, row_major> knn_graph,
                      experimental::nn_descent::index_params build_params)
 {
-  auto nn_descent_idx = experimental::nn_descent::index<float, IdxT>(res, knn_graph);
+  auto nn_descent_idx = experimental::nn_descent::index<IdxT>(res, knn_graph);
   experimental::nn_descent::build<DataT, IdxT>(res, build_params, dataset, nn_descent_idx);
 
   using internal_IdxT = typename std::make_unsigned<IdxT>::type;

@@ -1391,7 +1391,7 @@ template <typename T,
 void build(raft::resources const& res,
            const index_params& params,
            mdspan<const T, matrix_extent<int64_t>, row_major, Accessor> dataset,
-           index<DistData_t, IdxT>& idx)
+           index<IdxT>& idx)
 {
   RAFT_EXPECTS(dataset.extent(0) < std::numeric_limits<int>::max() - 1,
                "The dataset size for GNND should be less than %d",
@@ -1435,12 +1435,19 @@ void build(raft::resources const& res,
 
   GNND<const T, int> nnd(res, build_config);
 
-  // auto distances_graph = raft::make_device_matrix<float, int64_t>(res, 0, 0)
-  nnd.build(dataset.data_handle(),
-            dataset.extent(0),
-            int_graph.data_handle(),
-            params.return_distances,
-            idx.distances().value_or(raft::make_device_matrix<float, int64_t>(res, 0, 0).view()).data_handle());
+  if (idx.distances().has_value() || !params.return_distances) {
+    nnd.build(dataset.data_handle(),
+              dataset.extent(0),
+              int_graph.data_handle(),
+              params.return_distances,
+              idx.distances()
+                .value_or(raft::make_device_matrix<float, int64_t>(res, 0, 0).view())
+                .data_handle());
+  } else {
+    RAFT_EXPECTS(false,
+                 "Distance view not allocated. Using return_distances set to true requires "
+                 "distance view to be allocated.");
+  }
 
 #pragma omp parallel for
   for (size_t i = 0; i < static_cast<size_t>(dataset.extent(0)); i++) {
@@ -1455,9 +1462,9 @@ template <typename T,
           typename IdxT = uint32_t,
           typename Accessor =
             host_device_accessor<std::experimental::default_accessor<T>, memory_type::host>>
-index<DistData_t, IdxT> build(raft::resources const& res,
-                              const index_params& params,
-                              mdspan<const T, matrix_extent<int64_t>, row_major, Accessor> dataset)
+index<IdxT> build(raft::resources const& res,
+                  const index_params& params,
+                  mdspan<const T, matrix_extent<int64_t>, row_major, Accessor> dataset)
 {
   size_t intermediate_degree = params.intermediate_graph_degree;
   size_t graph_degree        = params.graph_degree;
@@ -1471,7 +1478,7 @@ index<DistData_t, IdxT> build(raft::resources const& res,
     graph_degree = intermediate_degree;
   }
 
-  index<DistData_t, IdxT> idx{
+  index<IdxT> idx{
     res, dataset.extent(0), static_cast<int64_t>(graph_degree), params.return_distances};
 
   build(res, params, dataset, idx);

diff --git a/cpp/include/raft/neighbors/nn_descent.cuh b/cpp/include/raft/neighbors/nn_descent.cuh
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -56,9 +56,9 @@ namespace raft::neighbors::experimental::nn_descent {
  * @return index<IdxT> index containing all-neighbors knn graph in host memory
  */
 template <typename T, typename IdxT = uint32_t>
-index<detail::DistData_t, IdxT> build(raft::resources const& res,
-                                      index_params const& params,
-                                      raft::device_matrix_view<const T, int64_t, row_major> dataset)
+index<IdxT> build(raft::resources const& res,
+                  index_params const& params,
+                  raft::device_matrix_view<const T, int64_t, row_major> dataset)
 {
   return detail::build<T, IdxT>(res, params, dataset);
 }
@@ -97,7 +97,7 @@ template <typename T, typename IdxT = uint32_t>
 void build(raft::resources const& res,
            index_params const& params,
            raft::device_matrix_view<const T, int64_t, row_major> dataset,
-           index<detail::DistData_t, IdxT>& idx)
+           index<IdxT>& idx)
 {
   detail::build<T, IdxT>(res, params, dataset, idx);
 }
@@ -130,9 +130,9 @@ void build(raft::resources const& res,
  * @return index<IdxT> index containing all-neighbors knn graph in host memory
  */
 template <typename T, typename IdxT = uint32_t>
-index<detail::DistData_t, IdxT> build(raft::resources const& res,
-                                      index_params const& params,
-                                      raft::host_matrix_view<const T, int64_t, row_major> dataset)
+index<IdxT> build(raft::resources const& res,
+                  index_params const& params,
+                  raft::host_matrix_view<const T, int64_t, row_major> dataset)
 {
   return detail::build<T, IdxT>(res, params, dataset);
 }
@@ -171,7 +171,7 @@ template <typename T, typename IdxT = uint32_t>
 void build(raft::resources const& res,
            index_params const& params,
            raft::host_matrix_view<const T, int64_t, row_major> dataset,
-           index<detail::DistData_t, IdxT>& idx)
+           index<IdxT>& idx)
 {
   detail::build<T, IdxT>(res, params, dataset, idx);
 }

@@ -69,10 +69,9 @@ struct index_params : ann::index_params {
  * International Conference on Information & Knowledge Management (CIKM '21). Association for
  * Computing Machinery, New York, NY, USA, 1929–1938. https://doi.org/10.1145/3459637.3482344
  *
- * @tparam T dtype to be used for constructing distances graph
  * @tparam IdxT dtype to be used for constructing knn-graph
  */
-template <typename T, typename IdxT>
+template <typename IdxT>
 struct index : ann::index {
  public:
   /**
@@ -96,7 +95,7 @@ struct index : ann::index {
       return_distances_(return_distances)
   {
     if (return_distances) {
-      distances_      = raft::make_device_matrix<T, int64_t>(res_, n_rows, n_cols);
+      distances_      = raft::make_device_matrix<float, int64_t>(res_, n_rows, n_cols);
       distances_view_ = distances_.value().view();
     }
   }
@@ -114,16 +113,16 @@ struct index : ann::index {
    * storing knn-graph distances
    * @param return_distances whether to allocate and get distances information
    */
-  index(
-    raft::resources const& res,
-    raft::host_matrix_view<IdxT, int64_t, raft::row_major> graph_view,
-    std::optional<raft::device_matrix_view<T, int64_t, row_major>> distances_view = std::nullopt,
-    bool return_distances                                                         = false)
+  index(raft::resources const& res,
+        raft::host_matrix_view<IdxT, int64_t, raft::row_major> graph_view,
+        std::optional<raft::device_matrix_view<float, int64_t, row_major>> distances_view =
+          std::nullopt,
+        bool return_distances = false)
     : ann::index(),
       res_{res},
       metric_{raft::distance::DistanceType::L2Expanded},
       graph_{raft::make_host_matrix<IdxT, int64_t, row_major>(0, 0)},
-      distances_{raft::make_device_matrix<T, int64_t>(res_, 0, 0)},
+      distances_{raft::make_device_matrix<float, int64_t>(res_, 0, 0)},
       graph_view_{graph_view},
       distances_view_(distances_view),
       return_distances_(return_distances)
@@ -156,7 +155,7 @@ struct index : ann::index {
 
   /** neighborhood graph distances [size, graph-degree] */
   [[nodiscard]] inline auto distances() noexcept
-    -> std::optional<device_matrix_view<T, int64_t, row_major>>
+    -> std::optional<device_matrix_view<float, int64_t, row_major>>
   {
     return distances_view_;
   }
@@ -172,10 +171,10 @@ struct index : ann::index {
   raft::resources const& res_;
   raft::distance::DistanceType metric_;
   raft::host_matrix<IdxT, int64_t, row_major> graph_;  // graph to return for non-int IdxT
-  std::optional<raft::device_matrix<T, int64_t, row_major>> distances_;
+  std::optional<raft::device_matrix<float, int64_t, row_major>> distances_;
   raft::host_matrix_view<IdxT, int64_t, row_major>
     graph_view_;  // view of graph for user provided matrix
-  std::optional<raft::device_matrix_view<T, int64_t, row_major>> distances_view_;
+  std::optional<raft::device_matrix_view<float, int64_t, row_major>> distances_view_;
   bool return_distances_;
 };