Using sparse public API functions from RAFT (#4389)

Authors:
  - Corey J. Nolet (https://github.com/cjnolet)

Approvers:
  - Divye Gala (https://github.com/divyegala)
  - Dante Gama Dessavre (https://github.com/dantegd)

URL: #4389

cpp/include/cuml/manifold/umap.hpp

@@ -16,16 +16,14 @@
 
 #pragma once
 
+#include <raft/sparse/coo.hpp>
+
 #include <cstddef>
 #include <cstdint>
 #include <memory>
 
 namespace raft {
 class handle_t;
-namespace sparse {
-template <typename T, typename Index_Type>
-class COO;
-};
 }  // namespace raft
 
 namespace ML {

cpp/src/dbscan/adjgraph/algo.cuh

@@ -23,7 +23,7 @@
 #include "pack.h"
 
 #include <raft/cuda_utils.cuh>
-#include <raft/sparse/convert/csr.cuh>
+#include <raft/sparse/convert/csr.hpp>
 
 using namespace thrust;
 
@@ -49,7 +49,7 @@ void launcher(const raft::handle_t& handle,
 
   exclusive_scan(handle.get_thrust_policy(), dev_vd, dev_vd + batch_size, dev_ex_scan);
 
-  raft::sparse::convert::csr_adj_graph_batched<Index_, TPB_X>(
+  raft::sparse::convert::csr_adj_graph_batched<Index_>(
     data.ex_scan, data.N, data.adjnnz, batch_size, data.adj, data.adj_graph, stream);
 
   CUDA_CHECK(cudaPeekAtLastError());

cpp/src/dbscan/runner.cuh

@@ -19,7 +19,7 @@
 #include <raft/cudart_utils.h>
 #include <common/nvtx.hpp>
 #include <label/classlabels.cuh>
-#include <raft/sparse/csr.cuh>
+#include <raft/sparse/csr.hpp>
 #include "adjgraph/runner.cuh"
 #include "corepoints/compute.cuh"
 #include "corepoints/exchange.cuh"
@@ -236,7 +236,7 @@ std::size_t run(const raft::handle_t& handle,
 
     CUML_LOG_DEBUG("--> Computing connected components");
     ML::PUSH_RANGE("Trace::Dbscan::WeakCC");
-    raft::sparse::weak_cc_batched<Index_, 1024>(
+    raft::sparse::weak_cc_batched<Index_>(
       i == 0 ? labels : labels_temp,
       ex_scan,
       adj_graph.data(),

cpp/src/dbscan/vertexdeg/precomputed.cuh

@@ -21,6 +21,7 @@
 #include <raft/cudart_utils.h>
 #include <cub/cub.cuh>
 #include <raft/cuda_utils.cuh>
+#include <raft/device_atomics.cuh>
 #include <raft/linalg/coalesced_reduction.cuh>
 #include <raft/linalg/reduce.cuh>
 

cpp/src/hdbscan/condensed_hierarchy.cu

@@ -24,8 +24,8 @@
 #include <rmm/device_uvector.hpp>
 #include <rmm/exec_policy.hpp>
 
-#include <raft/sparse/op/sort.h>
-#include <raft/sparse/convert/csr.cuh>
+#include <raft/sparse/convert/csr.hpp>
+#include <raft/sparse/op/sort.hpp>
 
 #include <thrust/execution_policy.h>
 #include <thrust/reduce.h>

cpp/src/hdbscan/detail/condense.cuh

@@ -25,8 +25,8 @@
 #include <rmm/device_uvector.hpp>
 #include <rmm/exec_policy.hpp>
 
-#include <raft/sparse/op/sort.h>
-#include <raft/sparse/convert/csr.cuh>
+#include <raft/sparse/convert/csr.hpp>
+#include <raft/sparse/op/sort.hpp>
 
 #include <cuml/cluster/hdbscan.hpp>
 

cpp/src/hdbscan/detail/extract.cuh

@@ -26,9 +26,9 @@
 #include <cuml/cluster/hdbscan.hpp>
 
 #include <raft/cudart_utils.h>
-#include <raft/sparse/op/sort.h>
 #include <raft/label/classlabels.cuh>
-#include <raft/sparse/convert/csr.cuh>
+#include <raft/sparse/convert/csr.hpp>
+#include <raft/sparse/op/sort.hpp>
 
 #include <rmm/device_uvector.hpp>
 #include <rmm/exec_policy.hpp>

cpp/src/hdbscan/detail/membership.cuh

@@ -23,8 +23,8 @@
 
 #include <raft/cudart_utils.h>
 
-#include <raft/sparse/op/sort.h>
-#include <raft/sparse/convert/csr.cuh>
+#include <raft/sparse/convert/csr.hpp>
+#include <raft/sparse/op/sort.hpp>
 
 #include <cuml/cluster/hdbscan.hpp>
 

cpp/src/hdbscan/detail/reachability.cuh

@@ -25,10 +25,10 @@
 
 #include <raft/linalg/unary_op.cuh>
 
-#include <raft/sparse/convert/csr.cuh>
+#include <raft/sparse/convert/csr.hpp>
 #include <raft/sparse/hierarchy/detail/connectivities.cuh>
-#include <raft/sparse/linalg/symmetrize.cuh>
-#include <raft/sparse/selection/knn_graph.cuh>
+#include <raft/sparse/linalg/symmetrize.hpp>
+#include <raft/sparse/selection/knn_graph.hpp>
 
 #include <rmm/device_uvector.hpp>
 #include <rmm/exec_policy.hpp>

cpp/src/hdbscan/detail/select.cuh

@@ -23,8 +23,8 @@
 
 #include <raft/cudart_utils.h>
 
-#include <raft/sparse/op/sort.h>
-#include <raft/sparse/convert/csr.cuh>
+#include <raft/sparse/convert/csr.hpp>
+#include <raft/sparse/op/sort.hpp>
 
 #include <cuml/cluster/hdbscan.hpp>
 

cpp/src/hdbscan/detail/stabilities.cuh

@@ -23,8 +23,8 @@
 
 #include <raft/cudart_utils.h>
 
-#include <raft/sparse/op/sort.h>
-#include <raft/sparse/convert/csr.cuh>
+#include <raft/sparse/convert/csr.hpp>
+#include <raft/sparse/op/sort.hpp>
 
 #include <cuml/cluster/hdbscan.hpp>
 

cpp/src/hdbscan/detail/utils.h

@@ -22,8 +22,8 @@
 
 #include <raft/cudart_utils.h>
 
-#include <raft/sparse/op/sort.h>
-#include <raft/sparse/convert/csr.cuh>
+#include <raft/sparse/convert/csr.hpp>
+#include <raft/sparse/op/sort.hpp>
 
 #include <cuml/cluster/hdbscan.hpp>
 

cpp/src/hdbscan/runner.h

@@ -23,7 +23,7 @@
 
 #include <cuml/common/logger.hpp>
 
-#include <raft/sparse/coo.cuh>
+#include <raft/sparse/coo.hpp>
 #include <raft/sparse/hierarchy/detail/agglomerative.cuh>
 #include <raft/sparse/hierarchy/detail/mst.cuh>
 

cpp/src/knn/knn_sparse.cu

@@ -17,7 +17,7 @@
 #include <cuml/common/logger.hpp>
 #include <cuml/neighbors/knn_sparse.hpp>
 
-#include <raft/sparse/selection/knn.cuh>
+#include <raft/sparse/selection/knn.hpp>
 
 #include <cusparse_v2.h>
 

cpp/src/spectral/spectral.cu

@@ -14,9 +14,9 @@
  * limitations under the License.
  */
 
-#include <raft/sparse/coo.cuh>
+#include <raft/sparse/coo.hpp>
 
-#include <raft/sparse/linalg/spectral.cuh>
+#include <raft/sparse/linalg/spectral.hpp>
 
 namespace raft {
 class handle_t;

cpp/src/tsne/distances.cuh

@@ -20,9 +20,9 @@
 #include <raft/linalg/distance_type.h>
 #include <cuml/neighbors/knn_sparse.hpp>
 #include <raft/linalg/eltwise.cuh>
-#include <raft/sparse/coo.cuh>
-#include <raft/sparse/linalg/symmetrize.cuh>
-#include <raft/sparse/selection/knn.cuh>
+#include <raft/sparse/coo.hpp>
+#include <raft/sparse/linalg/symmetrize.hpp>
+#include <raft/sparse/selection/knn.hpp>
 #include <selection/knn.cuh>
 
 #include <cuml/manifold/common.hpp>

cpp/src/tsne/utils.cuh

@@ -182,9 +182,9 @@ __global__ void min_max_kernel(
  * CUDA kernel to compute KL divergence
  */
 template <typename value_idx, typename value_t>
-__global__ void compute_kl_div_k(const value_t* restrict Ps,
+__global__ void compute_kl_div_k(const value_t* Ps,
                                  const value_t* Qs,
-                                 value_t* KL_divs,
+                                 value_t* __restrict__ KL_divs,
                                  const value_idx NNZ)
 {
   const auto index = (blockIdx.x * blockDim.x) + threadIdx.x;
@@ -200,7 +200,7 @@ __global__ void compute_kl_div_k(const value_t* restrict Ps,
  */
 template <typename value_t>
 value_t compute_kl_div(
-  value_t* restrict Ps, value_t* Qs, value_t* KL_divs, const size_t NNZ, cudaStream_t stream)
+  value_t* __restrict__ Ps, value_t* Qs, value_t* KL_divs, const size_t NNZ, cudaStream_t stream)
 {
   value_t P_sum = thrust::reduce(rmm::exec_policy(stream), Ps, Ps + NNZ);
   raft::linalg::scalarMultiply(Ps, Ps, 1.0f / P_sum, NNZ, stream);

cpp/src/umap/fuzzy_simpl_set/naive.cuh

@@ -23,9 +23,9 @@
 #include <raft/cudart_utils.h>
 #include <raft/cuda_utils.cuh>
 
-#include <raft/sparse/op/sort.h>
-#include <raft/sparse/coo.cuh>
-#include <raft/sparse/linalg/symmetrize.cuh>
+#include <raft/sparse/coo.hpp>
+#include <raft/sparse/linalg/symmetrize.hpp>
+#include <raft/sparse/op/sort.hpp>
 #include <raft/stats/mean.hpp>
 
 #include <cuda_runtime.h>
@@ -353,7 +353,7 @@ void launcher(int n,
    * one via a fuzzy union. (Symmetrize knn graph).
    */
   float set_op_mix_ratio = params->set_op_mix_ratio;
-  raft::sparse::linalg::coo_symmetrize<TPB_X, value_t>(
+  raft::sparse::linalg::coo_symmetrize<value_t>(
     &in,
     out,
     [set_op_mix_ratio] __device__(int row, int col, value_t result, value_t transpose) {

cpp/src/umap/fuzzy_simpl_set/runner.cuh

@@ -18,7 +18,7 @@
 #include <cuml/manifold/umapparams.h>
 #include "naive.cuh"
 
-#include <raft/sparse/coo.cuh>
+#include <raft/sparse/coo.hpp>
 
 namespace UMAPAlgo {
 

cpp/src/umap/init_embed/runner.cuh

@@ -18,7 +18,7 @@
 
 #include <cuml/manifold/umapparams.h>
 
-#include <raft/sparse/coo.cuh>
+#include <raft/sparse/coo.hpp>
 
 #include "random_algo.cuh"
 #include "spectral_algo.cuh"

cpp/src/umap/init_embed/spectral_algo.cuh

@@ -18,9 +18,8 @@
 
 #include <cuml/manifold/umapparams.h>
 
-#include <raft/sparse/coo.cuh>
-
 #include <raft/linalg/add.cuh>
+#include <raft/sparse/coo.hpp>
 
 #include <raft/linalg/transpose.h>
 #include <raft/random/rng.hpp>

cpp/src/umap/knn_graph/algo.cuh

@@ -22,7 +22,7 @@
 #include <cuml/neighbors/knn_sparse.hpp>
 #include <iostream>
 #include <raft/linalg/unary_op.cuh>
-#include <raft/sparse/selection/knn.cuh>
+#include <raft/sparse/selection/knn.hpp>
 #include <selection/knn.cuh>
 
 #include <raft/cudart_utils.h>

cpp/src/umap/runner.cuh

@@ -36,11 +36,11 @@
 #include <thrust/scan.h>
 #include <thrust/system/cuda/execution_policy.h>
 
-#include <raft/sparse/op/sort.h>
-#include <raft/sparse/convert/csr.cuh>
-#include <raft/sparse/coo.cuh>
-#include <raft/sparse/linalg/norm.cuh>
-#include <raft/sparse/op/filter.cuh>
+#include <raft/sparse/convert/csr.hpp>
+#include <raft/sparse/coo.hpp>
+#include <raft/sparse/linalg/norm.hpp>
+#include <raft/sparse/op/filter.hpp>
+#include <raft/sparse/op/sort.hpp>
 
 #include <raft/cuda_utils.cuh>
 
@@ -139,7 +139,7 @@ void _fit(const raft::handle_t& handle,
    * Remove zeros from simplicial set
    */
   raft::sparse::COO<value_t> cgraph_coo(stream);
-  raft::sparse::op::coo_remove_zeros<TPB_X, value_t>(&rgraph_coo, &cgraph_coo, stream);
+  raft::sparse::op::coo_remove_zeros<value_t>(&rgraph_coo, &cgraph_coo, stream);
   ML::POP_RANGE();
 
   /**
@@ -218,7 +218,7 @@ void _get_graph(const raft::handle_t& handle,
   /**
    * Remove zeros from simplicial set
    */
-  raft::sparse::op::coo_remove_zeros<TPB_X, value_t>(&rgraph_coo, cgraph_coo, stream);
+  raft::sparse::op::coo_remove_zeros<value_t>(&rgraph_coo, cgraph_coo, stream);
   ML::POP_RANGE();
 }
 
@@ -285,7 +285,7 @@ void _get_graph_supervised(
                                                 stream);
   CUDA_CHECK(cudaPeekAtLastError());
 
-  raft::sparse::op::coo_remove_zeros<TPB_X, value_t>(&tmp_coo, &rgraph_coo, stream);
+  raft::sparse::op::coo_remove_zeros<value_t>(&tmp_coo, &rgraph_coo, stream);
 
   /**
    * If target metric is 'categorical', perform
@@ -311,7 +311,7 @@ void _get_graph_supervised(
   raft::sparse::op::coo_sort<value_t>(cgraph_coo, stream);
 
   raft::sparse::COO<value_t> ocoo(stream);
-  raft::sparse::op::coo_remove_zeros<TPB_X, value_t>(cgraph_coo, &ocoo, stream);
+  raft::sparse::op::coo_remove_zeros<value_t>(cgraph_coo, &ocoo, stream);
   ML::POP_RANGE();
 }
 
@@ -389,7 +389,7 @@ void _fit_supervised(const raft::handle_t& handle,
                                                 stream);
   CUDA_CHECK(cudaPeekAtLastError());
 
-  raft::sparse::op::coo_remove_zeros<TPB_X, value_t>(&tmp_coo, &rgraph_coo, stream);
+  raft::sparse::op::coo_remove_zeros<value_t>(&tmp_coo, &rgraph_coo, stream);
 
   raft::sparse::COO<value_t> final_coo(stream);
 
@@ -417,7 +417,7 @@ void _fit_supervised(const raft::handle_t& handle,
   raft::sparse::op::coo_sort<value_t>(&final_coo, stream);
 
   raft::sparse::COO<value_t> ocoo(stream);
-  raft::sparse::op::coo_remove_zeros<TPB_X, value_t>(&final_coo, &ocoo, stream);
+  raft::sparse::op::coo_remove_zeros<value_t>(&final_coo, &ocoo, stream);
   ML::POP_RANGE();
 
   /**
@@ -551,14 +551,14 @@ void _transform(const raft::handle_t& handle,
   rmm::device_uvector<int> ia(inputs.n, stream);
 
   raft::sparse::convert::sorted_coo_to_csr(&graph_coo, row_ind.data(), stream);
-  raft::sparse::linalg::coo_degree<TPB_X>(&graph_coo, ia.data(), stream);
+  raft::sparse::linalg::coo_degree(&graph_coo, ia.data(), stream);
 
   rmm::device_uvector<value_t> vals_normed(graph_coo.nnz, stream);
   CUDA_CHECK(cudaMemsetAsync(vals_normed.data(), 0, graph_coo.nnz * sizeof(value_t), stream));
 
   CUML_LOG_DEBUG("Performing L1 normalization");
 
-  raft::sparse::linalg::csr_row_normalize_l1<TPB_X, value_t>(
+  raft::sparse::linalg::csr_row_normalize_l1<value_t>(
     row_ind.data(), graph_coo.vals(), graph_coo.nnz, graph_coo.n_rows, vals_normed.data(), stream);
 
   init_transform<TPB_X, value_t><<<grid_n, blk, 0, stream>>>(graph_coo.cols(),
@@ -612,7 +612,7 @@ void _transform(const raft::handle_t& handle,
    * Remove zeros
    */
   raft::sparse::COO<value_t> comp_coo(stream);
-  raft::sparse::op::coo_remove_zeros<TPB_X, value_t>(&graph_coo, &comp_coo, stream);
+  raft::sparse::op::coo_remove_zeros<value_t>(&graph_coo, &comp_coo, stream);
 
   ML::PUSH_RANGE("umap::optimization");
   CUML_LOG_DEBUG("Computing # of epochs for training each sample");