[WIP] Sparse semiring cleanup + hash table and batching strategies

cpp/CMakeLists.txt

@@ -211,7 +211,9 @@ set(RAFT_INCLUDE_DIRECTORIES
   ${RAFT_INCLUDE_DIR}
   ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}
   ${FAISS_INCLUDE_DIRS}
-  ${RMM_INCLUDE_DIRS})
+  ${RMM_INCLUDE_DIRS}
+  ${CUDA_CXX_DIR}/include
+  ${CUCO_DIR}/include)
 
 if(NOT CUB_IS_PART_OF_CTK)
   list(APPEND RAFT_INCLUDE_DIRECTORIES ${CUB_DIR}/src/cub)
@@ -290,7 +292,6 @@ if(BUILD_RAFT_TESTS)
     test/sparse/csr_transpose.cu
     test/sparse/degree.cu
     test/sparse/dist_coo_spmv.cu
-    test/sparse/dist_csr_spmv.cu
     test/sparse/distance.cu
     test/sparse/filter.cu
     test/sparse/knn.cu

cpp/cmake/Dependencies.cmake

@@ -16,6 +16,36 @@
 
 include(ExternalProject)
 
+##############################################################################
+# - libcudacxx
+
+set(CUDA_CXX_DIR ${CMAKE_CURRENT_BINARY_DIR}/libcudacxx CACHE STRING "Path to libcudacxx repo")
+
+ExternalProject_Add(libcudacxx
+  GIT_REPOSITORY    https://github.com/NVIDIA/libcudacxx.git
+  GIT_TAG           main
+  PREFIX            ${CUDA_CXX_DIR}
+  CONFIGURE_COMMAND ""
+  BUILD_COMMAND     ""
+  INSTALL_COMMAND   "")
+
+  set(CUDA_CXX_DIR ${CUDA_CXX_DIR}/src/libcudacxx/)
+
+##############################################################################
+# - cucollections - (header only) -----------------------------------------------------
+
+set(CUCO_DIR ${CMAKE_CURRENT_BINARY_DIR}/cuCollections CACHE STRING "Path to cuCollections repo")
+
+ExternalProject_Add(cuCollections
+  GIT_REPOSITORY    https://github.com/divyegala/cuCollections.git
+  GIT_TAG           view-initialization-ctor
+  PREFIX            ${CUCO_DIR}
+  CONFIGURE_COMMAND ""
+  BUILD_COMMAND     ""
+  INSTALL_COMMAND   "")
+
+  set(CUCO_DIR ${CUCO_DIR}/src/cuCollections/)
+
 ##############################################################################
 # - cub - (header only) ------------------------------------------------------
 
@@ -108,4 +138,6 @@ endif(BUILD_GTEST)
 if(NOT CUB_IS_PART_OF_CTK)
   add_dependencies(GTest::GTest cub)
 endif(NOT CUB_IS_PART_OF_CTK)
+add_dependencies(cuCollections libcudacxx)
+add_dependencies(GTest::GTest cuCollections)
 add_dependencies(FAISS::FAISS faiss)

cpp/include/raft/cudart_utils.h

@@ -267,12 +267,12 @@ void allocate(Type*& ptr, size_t len, bool setZero = false) {
 }
 
 /** helper method to get max usable shared mem per block parameter */
-inline int getSharedMemPerBlock() {
+__host__ __device__ inline int getSharedMemPerBlock() {
   int devId;
-  CUDA_CHECK(cudaGetDevice(&devId));
+  cudaGetDevice(&devId);
   int smemPerBlk;
-  CUDA_CHECK(cudaDeviceGetAttribute(&smemPerBlk,
-                                    cudaDevAttrMaxSharedMemoryPerBlock, devId));
+  cudaDeviceGetAttribute(&smemPerBlk, cudaDevAttrMaxSharedMemoryPerBlock,
+                         devId);
   return smemPerBlk;
 }
 

cpp/include/raft/linalg/distance_type.h

@@ -52,10 +52,16 @@ enum DistanceType : unsigned short {
   Haversine = 13,
   /** Bray-Curtis distance **/
   BrayCurtis = 14,
-  /** Jensen-Shannon distance **/
+  /** Jensen-Shannon distance**/
   JensenShannon = 15,
+  /** Hamming distance **/
+  HammingUnexpanded = 16,
+  /** KLDivergence **/
+  KLDivergence = 17,
+  /** RusselRao **/
+  RusselRaoExpanded = 18,
   /** Dice-Sorensen distance **/
-  DiceExpanded = 16,
+  DiceExpanded = 19,
   /** Precomputed (special value) **/
   Precomputed = 100
 };

cpp/include/raft/sparse/distance/bin_distance.cuh

@@ -85,7 +85,6 @@ void compute_bin_distance(value_t *out, const value_idx *Q_coo_rows,
                           const value_t *Q_data, value_idx Q_nnz,
                           const value_idx *R_coo_rows, const value_t *R_data,
                           value_idx R_nnz, value_idx m, value_idx n,
-                          cusparseHandle_t handle,
                           std::shared_ptr<raft::mr::device::allocator> alloc,
                           cudaStream_t stream, expansion_f expansion_func) {
   raft::mr::device::buffer<value_t> Q_norms(alloc, stream, m);
@@ -114,7 +113,8 @@ class jaccard_expanded_distances_t : public distances_t<value_t> {
   explicit jaccard_expanded_distances_t(
     const distances_config_t<value_idx, value_t> &config)
     : config_(&config),
-      workspace(config.allocator, config.stream, 0),
+      workspace(config.handle.get_device_allocator(),
+                config.handle.get_stream(), 0),
       ip_dists(config) {}
 
   void compute(value_t *out_dists) {
@@ -124,15 +124,16 @@ class jaccard_expanded_distances_t : public distances_t<value_t> {
     value_t *b_data = ip_dists.b_data_coo();
 
     raft::mr::device::buffer<value_idx> search_coo_rows(
-      config_->allocator, config_->stream, config_->a_nnz);
+      config_->handle.get_device_allocator(), config_->handle.get_stream(),
+      config_->a_nnz);
     raft::sparse::convert::csr_to_coo(config_->a_indptr, config_->a_nrows,
                                       search_coo_rows.data(), config_->a_nnz,
-                                      config_->stream);
+                                      config_->handle.get_stream());
 
     compute_bin_distance(
       out_dists, search_coo_rows.data(), config_->a_data, config_->a_nnz,
       b_indices, b_data, config_->b_nnz, config_->a_nrows, config_->b_nrows,
-      config_->handle, config_->allocator, config_->stream,
+      config_->handle.get_device_allocator(), config_->handle.get_stream(),
       [] __device__ __host__(value_t dot, value_t q_norm, value_t r_norm) {
         value_t q_r_union = q_norm + r_norm;
         value_t denom = q_r_union - dot;
@@ -163,7 +164,8 @@ class dice_expanded_distances_t : public distances_t<value_t> {
   explicit dice_expanded_distances_t(
     const distances_config_t<value_idx, value_t> &config)
     : config_(&config),
-      workspace(config.allocator, config.stream, 0),
+      workspace(config.handle.get_device_allocator(),
+                config.handle.get_stream(), 0),
       ip_dists(config) {}
 
   void compute(value_t *out_dists) {
@@ -173,15 +175,16 @@ class dice_expanded_distances_t : public distances_t<value_t> {
     value_t *b_data = ip_dists.b_data_coo();
 
     raft::mr::device::buffer<value_idx> search_coo_rows(
-      config_->allocator, config_->stream, config_->a_nnz);
+      config_->handle.get_device_allocator(), config_->handle.get_stream(),
+      config_->a_nnz);
     raft::sparse::convert::csr_to_coo(config_->a_indptr, config_->a_nrows,
                                       search_coo_rows.data(), config_->a_nnz,
-                                      config_->stream);
+                                      config_->handle.get_stream());
 
     compute_bin_distance(
       out_dists, search_coo_rows.data(), config_->a_data, config_->a_nnz,
       b_indices, b_data, config_->b_nnz, config_->a_nrows, config_->b_nrows,
-      config_->handle, config_->allocator, config_->stream,
+      config_->handle.get_device_allocator(), config_->handle.get_stream(),
       [] __device__ __host__(value_t dot, value_t q_norm, value_t r_norm) {
         value_t q_r_union = q_norm + r_norm;
         value_t dice = (2 * dot) / q_r_union;

cpp/include/raft/sparse/distance/common.h

@@ -16,15 +16,16 @@
 
 #pragma once
 
-#include <cusparse_v2.h>
-#include <raft/mr/device/allocator.hpp>
+#include <raft/handle.hpp>
 
 namespace raft {
 namespace sparse {
 namespace distance {
 
 template <typename value_idx, typename value_t>
 struct distances_config_t {
+  distances_config_t(raft::handle_t &handle_) : handle(handle_) {}
+
   // left side
   value_idx a_nrows;
   value_idx a_ncols;
@@ -41,10 +42,7 @@ struct distances_config_t {
   value_idx *b_indices;
   value_t *b_data;
 
-  cusparseHandle_t handle;
-
-  std::shared_ptr<raft::mr::device::allocator> allocator;
-  cudaStream_t stream;
+  raft::handle_t &handle;
 };
 
 template <typename value_t>