work through updates

cpp/src/neighbors/detail/knn_brute_force.cuh

@@ -266,11 +266,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/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/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/partition.cuh

@@ -45,17 +45,21 @@ 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> 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,
   weight_t* eigVals,
   weight_t* eigVecs)
 {
-  return cuvs::spectral::detail::partition<vertex_t, weight_t, EigenSolver, ClusterSolver>(
+  return cuvs::spectral::detail::partition<vertex_t, weight_t, nnz_t, EigenSolver, ClusterSolver>(
     handle, csr_m, eigen_solver, cluster_solver, clusters, eigVals, eigVecs);
 }
 

cpp/src/sparse/neighbors/detail/cross_component_nn.cuh

@@ -449,10 +449,10 @@ void min_components_by_color(raft::sparse::COO<value_t, value_idx>& coo,
  * is done
  * @param[in] metric distance metric
  */
-template <typename value_idx, typename value_t, typename red_op>
+template <typename value_idx, typename value_t, typename red_op, typename nnz_t = size_t>
 void cross_component_nn(
   raft::resources const& handle,
-  raft::sparse::COO<value_t, value_idx>& out,
+  raft::sparse::COO<value_t, value_idx, nnz_t>& out,
   const value_t* X,
   const value_idx* orig_colors,
   size_t n_rows,
@@ -510,7 +510,7 @@ void cross_component_nn(
    * Take the min for any duplicate colors
    */
   // Compute mask of duplicates
-  rmm::device_uvector<value_idx> out_index(n_rows + 1, stream);
+  rmm::device_uvector<nnz_t> out_index(n_rows + 1, stream);
   raft::sparse::op::compute_duplicates_mask(
     out_index.data(), colors.data(), nn_colors.data(), n_rows, stream);
 
@@ -520,13 +520,13 @@ void cross_component_nn(
                          out_index.data());
 
   // compute final size
-  value_idx size = 0;
+  nnz_t size = 0;
   raft::update_host(&size, out_index.data() + (out_index.size() - 1), 1, stream);
   raft::resource::sync_stream(handle, stream);
 
   size++;
 
-  raft::sparse::COO<value_t, value_idx> min_edges(stream);
+  raft::sparse::COO<value_t, value_idx, nnz_t> min_edges(stream);
   min_edges.allocate(size, n_rows, n_rows, true, stream);
 
   min_components_by_color(

cpp/tests/sparse/cluster/eigen_solvers.cu

@@ -36,17 +36,18 @@ TEST(Raft, EigenSolvers)
   using namespace raft::spectral::matrix;
   using index_type = int;
   using value_type = double;
+  using nnz_type   = int;
 
   raft::resources h;
   ASSERT_EQ(0, raft::resource::get_device_id(h));
 
   index_type* ro{nullptr};
   index_type* ci{nullptr};
   value_type* vs{nullptr};
-  index_type nnz   = 0;
+  nnz_type nnz     = 0;
   index_type nrows = 0;
 
-  sparse_matrix_t<index_type, value_type> sm1{h, ro, ci, vs, nrows, nnz};
+  sparse_matrix_t<index_type, value_type, nnz_type> sm1{h, ro, ci, vs, nrows, nnz};
   ASSERT_EQ(nullptr, sm1.row_offsets_);
 
   index_type neigvs{10};

cpp/tests/sparse/cluster/spectral.cu

@@ -73,13 +73,13 @@ TEST(Raft, Spectral)
   raft::update_device(indices.data(), h_indices.data(), h_indices.size(), handle.get_stream());
   raft::update_device(values.data(), h_values.data(), h_values.size(), handle.get_stream());
 
-  raft::spectral::matrix::sparse_matrix_t<int32_t, float> const matrix{
+  raft::spectral::matrix::sparse_matrix_t<int32_t, float, uint64_t> const matrix{
     handle,
     offsets.data(),
     indices.data(),
     values.data(),
     static_cast<int32_t>(offsets.size() - 1),
-    static_cast<int32_t>(indices.size())};
+    static_cast<uint64_t>(indices.size())};
 
   cuvs::spectral::eigen_solver_config_t<int32_t, float> eig_cfg{
     n_eigenvectors, evs_max_it, restartIter_lanczos, evs_tol, reorthog, seed1};