Account for RAFT sparse types updates

cpp/cmake/thirdparty/get_raft.cmake

@@ -61,8 +61,8 @@ endfunction()
 # To use a different RAFT locally, set the CMake variable
 # CPM_raft_SOURCE=/path/to/local/raft
 find_and_configure_raft(VERSION  ${RAFT_VERSION}.00
-        FORK                     ${RAFT_FORK}
-        PINNED_TAG               ${RAFT_PINNED_TAG}
+        FORK                     viclafargue
+        PINNED_TAG               fix-sparse-utilities
         ENABLE_MNMG_DEPENDENCIES OFF
         ENABLE_NVTX              OFF
         USE_RAFT_STATIC ${CUVS_USE_RAFT_STATIC}

cpp/src/cluster/detail/connectivities.cuh

@@ -43,8 +43,8 @@ template <Linkage dist_type, typename value_idx, typename value_t>
 struct distance_graph_impl {
   void run(raft::resources const& handle,
            const value_t* X,
-           size_t m,
-           size_t n,
+           value_idx m,
+           value_idx n,
            cuvs::distance::DistanceType metric,
            rmm::device_uvector<value_idx>& indptr,
            rmm::device_uvector<value_idx>& indices,
@@ -61,8 +61,8 @@ template <typename value_idx, typename value_t>
 struct distance_graph_impl<Linkage::KNN_GRAPH, value_idx, value_t> {
   void run(raft::resources const& handle,
            const value_t* X,
-           size_t m,
-           size_t n,
+           value_idx m,
+           value_idx n,
            cuvs::distance::DistanceType metric,
            rmm::device_uvector<value_idx>& indptr,
            rmm::device_uvector<value_idx>& indices,
@@ -131,8 +131,8 @@ RAFT_KERNEL fill_indices2(value_idx* indices, size_t m, size_t nnz)
 template <typename value_idx, typename value_t>
 void pairwise_distances(const raft::resources& handle,
                         const value_t* X,
-                        size_t m,
-                        size_t n,
+                        value_idx m,
+                        value_idx n,
                         cuvs::distance::DistanceType metric,
                         value_idx* indptr,
                         value_idx* indices,
@@ -183,8 +183,8 @@ template <typename value_idx, typename value_t>
 struct distance_graph_impl<Linkage::PAIRWISE, value_idx, value_t> {
   void run(const raft::resources& handle,
            const value_t* X,
-           size_t m,
-           size_t n,
+           value_idx m,
+           value_idx n,
            cuvs::distance::DistanceType metric,
            rmm::device_uvector<value_idx>& indptr,
            rmm::device_uvector<value_idx>& indices,
@@ -221,8 +221,8 @@ struct distance_graph_impl<Linkage::PAIRWISE, value_idx, value_t> {
 template <typename value_idx, typename value_t, Linkage dist_type>
 void get_distance_graph(raft::resources const& handle,
                         const value_t* X,
-                        size_t m,
-                        size_t n,
+                        value_idx m,
+                        value_idx n,
                         cuvs::distance::DistanceType metric,
                         rmm::device_uvector<value_idx>& indptr,
                         rmm::device_uvector<value_idx>& indices,

cpp/src/cluster/single_linkage.cuh

@@ -50,8 +50,8 @@ namespace cuvs::cluster::agglomerative {
 template <typename value_idx, typename value_t, Linkage dist_type = Linkage::KNN_GRAPH>
 void single_linkage(raft::resources const& handle,
                     const value_t* X,
-                    size_t m,
-                    size_t n,
+                    value_idx m,
+                    value_idx n,
                     cuvs::distance::DistanceType metric,
                     single_linkage_output<value_idx>* out,
                     int c,
@@ -94,8 +94,8 @@ void single_linkage(raft::resources const& handle,
 
   single_linkage<idx_t, value_t, dist_type>(handle,
                                             X.data_handle(),
-                                            static_cast<std::size_t>(X.extent(0)),
-                                            static_cast<std::size_t>(X.extent(1)),
+                                            X.extent(0),
+                                            X.extent(1),
                                             metric,
                                             &out_arrs,
                                             c.has_value() ? c.value() : DEFAULT_CONST_C,

cpp/src/neighbors/detail/knn_brute_force.cuh

@@ -273,11 +273,13 @@ void tiled_brute_force_knn(const raft::resources& handle,
                          });
       }
 
+      std::optional<raft::device_matrix_view<const IndexType, int64_t, raft::row_major>>
+        null_optional = std::nullopt;
       cuvs::selection::select_k(
         handle,
         raft::make_device_matrix_view<const DistanceT, int64_t, raft::row_major>(
           temp_distances.data(), current_query_size, current_centroid_size),
-        std::nullopt,
+        null_optional,
         raft::make_device_matrix_view<DistanceT, int64_t, raft::row_major>(
           distances + i * k, current_query_size, current_k),
         raft::make_device_matrix_view<IndexType, int64_t, raft::row_major>(

cpp/src/neighbors/detail/reachability.cuh

@@ -186,17 +186,17 @@ void mutual_reachability_knn_l2(const raft::resources& handle,
       epilogue);
 }
 
-template <typename value_idx, typename value_t>
+template <typename value_idx, typename value_t, typename nnz_t>
 void mutual_reachability_graph(const raft::resources& handle,
                                const value_t* X,
-                               size_t m,
-                               size_t n,
+                               value_idx m,
+                               value_idx n,
                                cuvs::distance::DistanceType metric,
                                int min_samples,
                                value_t alpha,
                                value_idx* indptr,
                                value_t* core_dists,
-                               raft::sparse::COO<value_t, value_idx>& out)
+                               raft::sparse::COO<value_t, value_idx, nnz_t>& out)
 {
   RAFT_EXPECTS(metric == cuvs::distance::DistanceType::L2SqrtExpanded,
                "Currently only L2 expanded distance is supported");
@@ -228,8 +228,14 @@ void mutual_reachability_graph(const raft::resources& handle,
                     coo_rows.data(),
                     [min_samples] __device__(value_idx c) -> value_idx { return c / min_samples; });
 
-  raft::sparse::linalg::symmetrize(
-    handle, coo_rows.data(), inds.data(), dists.data(), m, m, min_samples * m, out);
+  raft::sparse::linalg::symmetrize(handle,
+                                   coo_rows.data(),
+                                   inds.data(),
+                                   dists.data(),
+                                   m,
+                                   m,
+                                   static_cast<nnz_t>(min_samples * m),
+                                   out);
 
   raft::sparse::convert::sorted_coo_to_csr(out.rows(), out.nnz, indptr, m + 1, stream);
 

cpp/src/neighbors/reachability.cu

@@ -25,7 +25,7 @@ void mutual_reachability_graph(const raft::resources& handle,
                                int min_samples,
                                raft::device_vector_view<int> indptr,
                                raft::device_vector_view<float> core_dists,
-                               raft::sparse::COO<float, int>& out,
+                               raft::sparse::COO<float, int, size_t>& out,
                                cuvs::distance::DistanceType metric,
                                float alpha)
 {
@@ -34,7 +34,7 @@ void mutual_reachability_graph(const raft::resources& handle,
   RAFT_EXPECTS(indptr.extent(0) == static_cast<size_t>(X.extent(0) + 1),
                "indptr doesn't have expected size");
 
-  cuvs::neighbors::detail::reachability::mutual_reachability_graph<int, float>(
+  cuvs::neighbors::detail::reachability::mutual_reachability_graph<int, float, size_t>(
     handle,
     X.data_handle(),
     X.extent(0),

cpp/src/sparse/cluster/detail/spectral.cuh

@@ -56,12 +56,13 @@ void fit_embedding(raft::resources const& handle,
    */
   using index_type = int;
   using value_type = T;
+  using nnz_type   = int;
 
   index_type* ro = src_offsets.data();
   index_type* ci = dst_cols.data();
   value_type* vs = dst_vals.data();
 
-  raft::spectral::matrix::sparse_matrix_t<index_type, value_type> const r_csr_m{
+  raft::spectral::matrix::sparse_matrix_t<index_type, value_type, nnz_type> const r_csr_m{
     handle, ro, ci, vs, n, nnz};
 
   index_type neigvs       = n_components + 1;

cpp/src/sparse/cluster/detail/spectral/modularity_maximization.hpp

@@ -66,10 +66,14 @@ namespace detail {
  *    performed.
  *  @return error flag.
  */
-template <typename vertex_t, typename weight_t, typename EigenSolver, typename ClusterSolver>
+template <typename vertex_t,
+          typename weight_t,
+          typename nnz_t,
+          typename EigenSolver,
+          typename ClusterSolver>
 std::tuple<vertex_t, weight_t, vertex_t> modularity_maximization(
   raft::resources const& handle,
-  raft::spectral::matrix::sparse_matrix_t<vertex_t, weight_t> const& csr_m,
+  raft::spectral::matrix::sparse_matrix_t<vertex_t, weight_t, nnz_t> const& csr_m,
   EigenSolver const& eigen_solver,
   ClusterSolver const& cluster_solver,
   vertex_t* __restrict__ clusters,
@@ -91,7 +95,7 @@ std::tuple<vertex_t, weight_t, vertex_t> modularity_maximization(
   // Compute eigenvectors of Modularity Matrix
 
   // Initialize Modularity Matrix
-  raft::spectral::matrix::modularity_matrix_t<vertex_t, weight_t> B{handle, csr_m};
+  raft::spectral::matrix::modularity_matrix_t<vertex_t, weight_t, nnz_t> B{handle, csr_m};
 
   auto eigen_config = eigen_solver.get_config();
   auto nEigVecs     = eigen_config.n_eigVecs;
@@ -127,12 +131,13 @@ std::tuple<vertex_t, weight_t, vertex_t> modularity_maximization(
  *  @param clusters (Input, device memory, n entries) Cluster assignments.
  *  @param modularity On exit, modularity
  */
-template <typename vertex_t, typename weight_t>
-void analyzeModularity(raft::resources const& handle,
-                       raft::spectral::matrix::sparse_matrix_t<vertex_t, weight_t> const& csr_m,
-                       vertex_t nClusters,
-                       vertex_t const* __restrict__ clusters,
-                       weight_t& modularity)
+template <typename vertex_t, typename weight_t, typename nnz_t>
+void analyzeModularity(
+  raft::resources const& handle,
+  raft::spectral::matrix::sparse_matrix_t<vertex_t, weight_t, nnz_t> const& csr_m,
+  vertex_t nClusters,
+  vertex_t const* __restrict__ clusters,
+  weight_t& modularity)
 {
   RAFT_EXPECTS(clusters != nullptr, "Null clusters buffer.");
 
@@ -152,7 +157,7 @@ void analyzeModularity(raft::resources const& handle,
     raft::linalg::detail::cublassetpointermode(cublas_h, CUBLAS_POINTER_MODE_HOST, stream));
 
   // Initialize Modularity
-  raft::spectral::matrix::modularity_matrix_t<vertex_t, weight_t> B{handle, csr_m};
+  raft::spectral::matrix::modularity_matrix_t<vertex_t, weight_t, nnz_t> B{handle, csr_m};
 
   // Initialize output
   modularity = 0;

cpp/src/sparse/cluster/detail/spectral/partition.hpp

@@ -66,10 +66,14 @@ namespace detail {
  *    performed.
  *  @return statistics: number of eigensolver iterations, .
  */
-template <typename vertex_t, typename weight_t, typename EigenSolver, typename ClusterSolver>
+template <typename vertex_t,
+          typename weight_t,
+          typename nnz_t,
+          typename EigenSolver,
+          typename ClusterSolver>
 std::tuple<vertex_t, weight_t, vertex_t> partition(
   raft::resources const& handle,
-  raft::spectral::matrix::sparse_matrix_t<vertex_t, weight_t> const& csr_m,
+  raft::spectral::matrix::sparse_matrix_t<vertex_t, weight_t, nnz_t> const& csr_m,
   EigenSolver const& eigen_solver,
   ClusterSolver const& cluster_solver,
   vertex_t* __restrict__ clusters,
@@ -97,7 +101,7 @@ std::tuple<vertex_t, weight_t, vertex_t> partition(
 
   // Initialize Laplacian
   /// sparse_matrix_t<vertex_t, weight_t> A{handle, graph};
-  raft::spectral::matrix::laplacian_matrix_t<vertex_t, weight_t> L{handle, csr_m};
+  raft::spectral::matrix::laplacian_matrix_t<vertex_t, weight_t, nnz_t> L{handle, csr_m};
 
   auto eigen_config = eigen_solver.get_config();
   auto nEigVecs     = eigen_config.n_eigVecs;
@@ -135,13 +139,14 @@ std::tuple<vertex_t, weight_t, vertex_t> partition(
  *  @param cost On exit, partition cost function.
  *  @return error flag.
  */
-template <typename vertex_t, typename weight_t>
-void analyzePartition(raft::resources const& handle,
-                      raft::spectral::matrix::sparse_matrix_t<vertex_t, weight_t> const& csr_m,
-                      vertex_t nClusters,
-                      const vertex_t* __restrict__ clusters,
-                      weight_t& edgeCut,
-                      weight_t& cost)
+template <typename vertex_t, typename weight_t, typename nnz_t>
+void analyzePartition(
+  raft::resources const& handle,
+  raft::spectral::matrix::sparse_matrix_t<vertex_t, weight_t, nnz_t> const& csr_m,
+  vertex_t nClusters,
+  const vertex_t* __restrict__ clusters,
+  weight_t& edgeCut,
+  weight_t& cost)
 {
   RAFT_EXPECTS(clusters != nullptr, "Null clusters buffer.");
 
@@ -163,7 +168,7 @@ void analyzePartition(raft::resources const& handle,
 
   // Initialize Laplacian
   /// sparse_matrix_t<vertex_t, weight_t> A{handle, graph};
-  raft::spectral::matrix::laplacian_matrix_t<vertex_t, weight_t> L{handle, csr_m};
+  raft::spectral::matrix::laplacian_matrix_t<vertex_t, weight_t, nnz_t> L{handle, csr_m};
 
   // Initialize output
   cost    = 0;

cpp/src/sparse/cluster/detail/spectral/spectral_util.cuh

@@ -136,16 +136,17 @@ struct equal_to_i_op {
 
 // Construct indicator vector for ith partition
 //
-template <typename vertex_t, typename edge_t, typename weight_t>
-bool construct_indicator(raft::resources const& handle,
-                         edge_t index,
-                         edge_t n,
-                         weight_t& clustersize,
-                         weight_t& partStats,
-                         vertex_t const* __restrict__ clusters,
-                         raft::spectral::matrix::vector_t<weight_t>& part_i,
-                         raft::spectral::matrix::vector_t<weight_t>& Bx,
-                         raft::spectral::matrix::laplacian_matrix_t<vertex_t, weight_t> const& B)
+template <typename vertex_t, typename edge_t, typename weight_t, typename nnz_t>
+bool construct_indicator(
+  raft::resources const& handle,
+  edge_t index,
+  edge_t n,
+  weight_t& clustersize,
+  weight_t& partStats,
+  vertex_t const* __restrict__ clusters,
+  raft::spectral::matrix::vector_t<weight_t>& part_i,
+  raft::spectral::matrix::vector_t<weight_t>& Bx,
+  raft::spectral::matrix::laplacian_matrix_t<vertex_t, weight_t, nnz_t> const& B)
 {
   auto stream             = raft::resource::get_cuda_stream(handle);
   auto cublas_h           = raft::resource::get_cublas_handle(handle);

cpp/src/sparse/cluster/eigen_solvers.cuh

@@ -51,7 +51,7 @@ struct lanczos_solver_t {
 
   index_type_t solve_smallest_eigenvectors(
     raft::resources const& handle,
-    raft::spectral::matrix::sparse_matrix_t<index_type_t, value_type_t> const& A,
+    raft::spectral::matrix::sparse_matrix_t<index_type_t, value_type_t, size_type_t> const& A,
     value_type_t* __restrict__ eigVals,
     value_type_t* __restrict__ eigVecs) const
   {
@@ -74,7 +74,7 @@ struct lanczos_solver_t {
 
   index_type_t solve_largest_eigenvectors(
     raft::resources const& handle,
-    raft::spectral::matrix::sparse_matrix_t<index_type_t, value_type_t> const& A,
+    raft::spectral::matrix::sparse_matrix_t<index_type_t, value_type_t, size_type_t> const& A,
     value_type_t* __restrict__ eigVals,
     value_type_t* __restrict__ eigVecs) const
   {

cpp/src/sparse/cluster/modularity_maximization.cuh

@@ -43,18 +43,22 @@ namespace spectral {
  *  @param eigVecs Output eigenvector array pointer on device
  *  @return statistics: number of eigensolver iterations, .
  */
-template <typename vertex_t, typename weight_t, typename EigenSolver, typename ClusterSolver>
+template <typename vertex_t,
+          typename weight_t,
+          typename nnz_t,
+          typename EigenSolver,
+          typename ClusterSolver>
 std::tuple<vertex_t, weight_t, vertex_t> modularity_maximization(
   raft::resources const& handle,
-  raft::spectral::matrix::sparse_matrix_t<vertex_t, weight_t> const& csr_m,
+  raft::spectral::matrix::sparse_matrix_t<vertex_t, weight_t, nnz_t> const& csr_m,
   EigenSolver const& eigen_solver,
   ClusterSolver const& cluster_solver,
   vertex_t* __restrict__ clusters,
   weight_t* eigVals,
   weight_t* eigVecs)
 {
   return cuvs::spectral::detail::
-    modularity_maximization<vertex_t, weight_t, EigenSolver, ClusterSolver>(
+    modularity_maximization<vertex_t, weight_t, nnz_t, EigenSolver, ClusterSolver>(
       handle, csr_m, eigen_solver, cluster_solver, clusters, eigVals, eigVecs);
 }
 //===================================================
@@ -69,14 +73,15 @@ std::tuple<vertex_t, weight_t, vertex_t> modularity_maximization(
  *  @param clusters (Input, device memory, n entries) Cluster assignments.
  *  @param modularity On exit, modularity
  */
-template <typename vertex_t, typename weight_t>
-void analyzeModularity(raft::resources const& handle,
-                       raft::spectral::matrix::sparse_matrix_t<vertex_t, weight_t> const& csr_m,
-                       vertex_t nClusters,
-                       vertex_t const* __restrict__ clusters,
-                       weight_t& modularity)
+template <typename vertex_t, typename weight_t, typename nnz_t>
+void analyzeModularity(
+  raft::resources const& handle,
+  raft::spectral::matrix::sparse_matrix_t<vertex_t, weight_t, nnz_t> const& csr_m,
+  vertex_t nClusters,
+  vertex_t const* __restrict__ clusters,
+  weight_t& modularity)
 {
-  cuvs::spectral::detail::analyzeModularity<vertex_t, weight_t>(
+  cuvs::spectral::detail::analyzeModularity<vertex_t, weight_t, nnz_t>(
     handle, csr_m, nClusters, clusters, modularity);
 }