Use RAFT template instantations for distances

ci/gpu/build.sh

@@ -184,6 +184,11 @@ else
     cd $LIBCUML_BUILD_DIR
     chrpath -d ./test/ml
     patchelf --replace-needed `patchelf --print-needed ./test/ml | grep faiss` libfaiss.so ./test/ml
+    cp _deps/raft-build/libraft_nn.so $PWD
+    patchelf --replace-needed `patchelf --print-needed libraft_nn.so | grep faiss` libfaiss.so libraft_nn.so
+    cp _deps/raft-build/libraft_distance.so $PWD
+
+    gpuci_logger "Running libcuml binaries"
     GTEST_OUTPUT="xml:${WORKSPACE}/test-results/libcuml_cpp/" ./test/ml
 
     CONDA_FILE=`find ${CONDA_ARTIFACT_PATH} -name "libcuml*.tar.bz2"`

cpp/CMakeLists.txt

@@ -397,6 +397,8 @@ if(BUILD_CUML_CPP_LIBRARY)
       rmm::rmm
       cuml::Thrust
       raft::raft
+      raft::raft_nn
+      raft::raft_distance
     PRIVATE
       CUDA::cublas
       CUDA::cufft

cpp/bench/CMakeLists.txt

@@ -45,6 +45,8 @@ if(BUILD_CUML_BENCH)
       cuml::${CUML_CPP_TARGET}
       benchmark::benchmark
       raft::raft
+      raft::raft_nn
+      raft::raft_distance
       $<IF:$<BOOL:${Treelite_ADDED}>,treelite::treelite_static,treelite::treelite>
       $<IF:$<BOOL:${Treelite_ADDED}>,treelite::treelite_runtime_static,treelite::treelite_runtime>
   )
@@ -87,6 +89,8 @@ if(BUILD_CUML_PRIMS_BENCH)
       CUDA::cublas
       benchmark::benchmark
       raft::raft
+      raft::raft_nn
+      raft::raft_distance
       $<IF:$<BOOL:${Treelite_ADDED}>,treelite::treelite_static,treelite::treelite>
       $<IF:$<BOOL:${Treelite_ADDED}>,treelite::treelite_runtime_static,treelite::treelite_runtime>
   )

cpp/bench/prims/distance_common.cuh

@@ -17,6 +17,7 @@
 #include <raft/cudart_utils.h>
 #include <common/ml_benchmark.hpp>
 #include <raft/distance/distance.hpp>
+#include <raft/distance/specializations.hpp>
 
 namespace MLCommon {
 namespace Bench {

cpp/cmake/thirdparty/get_raft.cmake

@@ -18,33 +18,33 @@ function(find_and_configure_raft)
 
     set(oneValueArgs VERSION FORK PINNED_TAG)
     cmake_parse_arguments(PKG "${options}" "${oneValueArgs}"
-                          "${multiValueArgs}" ${ARGN} )
+            "${multiValueArgs}" ${ARGN} )
 
     if(DEFINED CPM_raft_SOURCE OR NOT DISABLE_FORCE_CLONE_RAFT)
-      set(CPM_DL_ALL_CACHE ${CPM_DOWNLOAD_ALL})
-      set(CPM_DOWNLOAD_ALL ON)
+        set(CPM_DL_ALL_CACHE ${CPM_DOWNLOAD_ALL})
+        set(CPM_DOWNLOAD_ALL ON)
     endif()
 
     rapids_cpm_find(raft ${PKG_VERSION}
-      GLOBAL_TARGETS      raft::raft
-      BUILD_EXPORT_SET    cuml-exports
-      INSTALL_EXPORT_SET  cuml-exports
-        CPM_ARGS
+            GLOBAL_TARGETS      raft::raft
+            BUILD_EXPORT_SET    cuml-exports
+            INSTALL_EXPORT_SET  cuml-exports
+            CPM_ARGS
             GIT_REPOSITORY https://github.com/${PKG_FORK}/raft.git
             GIT_TAG        ${PKG_PINNED_TAG}
             SOURCE_SUBDIR  cpp
             OPTIONS
-              "BUILD_TESTS OFF"
-    )
+            "BUILD_TESTS OFF"
+            )
 
     if(raft_ADDED)
-      message(VERBOSE "CUML: Using RAFT located in ${raft_SOURCE_DIR}")
+        message(VERBOSE "CUML: Using RAFT located in ${raft_SOURCE_DIR}")
     else()
-      message(VERBOSE "CUML: Using RAFT located in ${raft_DIR}")
+        message(VERBOSE "CUML: Using RAFT located in ${raft_DIR}")
     endif()
 
     if(DEFINED CPM_raft_SOURCE OR NOT DISABLE_FORCE_CLONE_RAFT)
-      set(CPM_DOWNLOAD_ALL ${CPM_DL_ALL_CACHE})
+        set(CPM_DOWNLOAD_ALL ${CPM_DL_ALL_CACHE})
     endif()
 
 endfunction()
@@ -56,6 +56,6 @@ set(CUML_BRANCH_VERSION_raft "${CUML_VERSION_MAJOR}.${CUML_VERSION_MINOR}")
 # To use a different RAFT locally, set the CMake variable
 # CPM_raft_SOURCE=/path/to/local/raft
 find_and_configure_raft(VERSION    ${CUML_MIN_VERSION_raft}
-                        FORK       rapidsai
-                        PINNED_TAG branch-${CUML_BRANCH_VERSION_raft}
-                        )
+        FORK       rapidsai
+        PINNED_TAG branch-${CUML_BRANCH_VERSION_raft}
+        )

cpp/src/knn/knn.cu

@@ -19,7 +19,9 @@
 #include <raft/label/classlabels.cuh>
 #include <raft/spatial/knn/ann.hpp>
 #include <raft/spatial/knn/ball_cover.hpp>
+
 #include <raft/spatial/knn/knn.hpp>
+#include <raft/spatial/knn/specializations.hpp>
 #include <rmm/device_uvector.hpp>
 
 #include <cuml/common/logger.hpp>
@@ -49,20 +51,20 @@ void brute_force_knn(const raft::handle_t& handle,
 {
   ASSERT(input.size() == sizes.size(), "input and sizes vectors must be the same size");
 
-  raft::spatial::knn::brute_force_knn(handle,
-                                      input,
-                                      sizes,
-                                      D,
-                                      search_items,
-                                      n,
-                                      res_I,
-                                      res_D,
-                                      k,
-                                      rowMajorIndex,
-                                      rowMajorQuery,
-                                      nullptr,
-                                      metric,
-                                      metric_arg);
+  raft::spatial::knn::brute_force_knn<int64_t, float, int>(handle,
+                                                           input,
+                                                           sizes,
+                                                           D,
+                                                           search_items,
+                                                           n,
+                                                           res_I,
+                                                           res_D,
+                                                           k,
+                                                           rowMajorIndex,
+                                                           rowMajorQuery,
+                                                           nullptr,
+                                                           metric,
+                                                           metric_arg);
 }
 
 void rbc_build_index(const raft::handle_t& handle,

cpp/src/knn/knn_opg_common.cuh

@@ -427,7 +427,7 @@ void perform_local_knn(opg_knn_param<in_t, ind_t, dist_t, out_t>& params,
                        size_t query_size)
 {
   std::vector<in_t*> ptrs(params.idx_data->size());
-  std::vector<int> sizes(params.idx_data->size());
+  std::vector<std::size_t> sizes(params.idx_data->size());
 
   for (std::size_t cur_idx = 0; cur_idx < params.idx_data->size(); cur_idx++) {
     ptrs[cur_idx]  = params.idx_data->at(cur_idx)->ptr;
@@ -444,19 +444,20 @@ void perform_local_knn(opg_knn_param<in_t, ind_t, dist_t, out_t>& params,
 
   // ID ranges need to be offset by each local partition's
   // starting indices.
-  raft::spatial::knn::brute_force_knn(handle,
-                                      ptrs,
-                                      sizes,
-                                      params.idx_desc->N,
-                                      query,
-                                      query_size,
-                                      work.res_I.data(),
-                                      work.res_D.data(),
-                                      params.k,
-                                      params.rowMajorIndex,
-                                      params.rowMajorQuery,
-                                      &start_indices_long,
-                                      raft::distance::DistanceType::L2SqrtExpanded);
+  raft::spatial::knn::brute_force_knn<std::int64_t, float, std::size_t>(
+    handle,
+    ptrs,
+    sizes,
+    params.idx_desc->N,
+    query,
+    query_size,
+    work.res_I.data(),
+    work.res_D.data(),
+    params.k,
+    params.rowMajorIndex,
+    params.rowMajorQuery,
+    &start_indices_long,
+    raft::distance::DistanceType::L2SqrtExpanded);
   CUDA_CHECK(cudaStreamSynchronize(handle.get_stream()));
   CUDA_CHECK(cudaPeekAtLastError());
 }

cpp/src/metrics/pairwise_distance.cu

@@ -18,6 +18,7 @@
 #include <raft/sparse/distance/common.h>
 #include <cuml/metrics/metrics.hpp>
 #include <raft/distance/distance.hpp>
+#include <raft/distance/specializations.hpp>
 #include <raft/handle.hpp>
 #include <raft/sparse/distance/distance.hpp>
 #include "pairwise_distance_canberra.cuh"

cpp/src/metrics/pairwise_distance_canberra.cuh

@@ -18,6 +18,7 @@
 #pragma once
 
 #include <raft/distance/distance.hpp>
+#include <raft/distance/specializations.hpp>
 #include <raft/handle.hpp>
 
 namespace ML {

cpp/src/metrics/pairwise_distance_chebyshev.cuh

@@ -15,7 +15,9 @@
  * limitations under the License.
  */
 #pragma once
+
 #include <raft/distance/distance.hpp>
+#include <raft/distance/specializations.hpp>
 #include <raft/handle.hpp>
 
 namespace ML {

cpp/src/metrics/pairwise_distance_correlation.cuh

@@ -18,6 +18,7 @@
 #pragma once
 
 #include <raft/distance/distance.hpp>
+#include <raft/distance/specializations.hpp>
 #include <raft/handle.hpp>
 
 namespace ML {

cpp/src/metrics/pairwise_distance_cosine.cuh

@@ -17,6 +17,7 @@
 #pragma once
 
 #include <raft/distance/distance.hpp>
+#include <raft/distance/specializations.hpp>
 #include <raft/handle.hpp>
 
 namespace ML {

cpp/src/metrics/pairwise_distance_euclidean.cuh

@@ -16,6 +16,7 @@
  */
 #pragma once
 #include <raft/distance/distance.hpp>
+#include <raft/distance/specializations.hpp>
 #include <raft/handle.hpp>
 
 namespace ML {

cpp/src/metrics/pairwise_distance_hamming.cuh

@@ -18,6 +18,7 @@
 #pragma once
 
 #include <raft/distance/distance.hpp>
+#include <raft/distance/specializations.hpp>
 #include <raft/handle.hpp>
 
 namespace ML {

cpp/src/metrics/pairwise_distance_hellinger.cuh

@@ -17,6 +17,7 @@
 #pragma once
 
 #include <raft/distance/distance.hpp>
+#include <raft/distance/specializations.hpp>
 #include <raft/handle.hpp>
 
 namespace ML {

cpp/src/metrics/pairwise_distance_jensen_shannon.cuh

@@ -18,6 +18,7 @@
 #pragma once
 
 #include <raft/distance/distance.hpp>
+#include <raft/distance/specializations.hpp>
 #include <raft/handle.hpp>
 
 namespace ML {

cpp/src/metrics/pairwise_distance_kl_divergence.cuh

@@ -18,6 +18,7 @@
 #pragma once
 
 #include <raft/distance/distance.hpp>
+#include <raft/distance/specializations.hpp>
 #include <raft/handle.hpp>
 
 namespace ML {

cpp/src/metrics/pairwise_distance_l1.cuh

@@ -15,7 +15,9 @@
  * limitations under the License.
  */
 #pragma once
+
 #include <raft/distance/distance.hpp>
+#include <raft/distance/specializations.hpp>
 #include <raft/handle.hpp>
 
 namespace ML {

cpp/src/metrics/pairwise_distance_minkowski.cuh

@@ -16,7 +16,9 @@
  */
 
 #pragma once
+
 #include <raft/distance/distance.hpp>
+#include <raft/distance/specializations.hpp>
 #include <raft/handle.hpp>
 
 namespace ML {

cpp/src/metrics/pairwise_distance_russell_rao.cuh

@@ -18,6 +18,7 @@
 #pragma once
 
 #include <raft/distance/distance.hpp>
+#include <raft/distance/specializations.hpp>
 #include <raft/handle.hpp>
 
 namespace ML {

cpp/src/metrics/trustworthiness.cu

@@ -19,6 +19,8 @@
 #include <cuml/metrics/metrics.hpp>
 
 #include <raft/distance/distance.hpp>
+#include <raft/distance/specializations.hpp>
+
 #include <raft/handle.hpp>
 
 namespace ML {

cpp/src/tsne/distances.cuh

@@ -66,19 +66,27 @@ void get_distances(const raft::handle_t& handle,
   std::vector<float*> input_vec = {input.X};
   std::vector<int> sizes_vec    = {input.n};
 
-  raft::spatial::knn::brute_force_knn(handle,
-                                      input_vec,
-                                      sizes_vec,
-                                      input.d,
-                                      input.X,
-                                      input.n,
-                                      k_graph.knn_indices,
-                                      k_graph.knn_dists,
-                                      k_graph.n_neighbors,
-                                      true,
-                                      true,
-                                      nullptr,
-                                      DEFAULT_DISTANCE_METRIC);
+  /**
+ * std::vector<float *> &input, std::vector<int> &sizes,
+                     IntType D, float *search_items, IntType n, int64_t *res_I,
+                     float *res_D, IntType k,
+                     std::shared_ptr<deviceAllocator> allocator,
+                     cudaStream_t userStream,
+ */
+
+  raft::spatial::knn::brute_force_knn<int64_t, float, int>(handle,
+                                                           input_vec,
+                                                           sizes_vec,
+                                                           input.d,
+                                                           input.X,
+                                                           input.n,
+                                                           k_graph.knn_indices,
+                                                           k_graph.knn_dists,
+                                                           k_graph.n_neighbors,
+                                                           true,
+                                                           true,
+                                                           nullptr,
+                                                           DEFAULT_DISTANCE_METRIC);
 }
 
 // dense, int32 indices

cpp/src_prims/metrics/trustworthiness_score.cuh

@@ -15,7 +15,9 @@
  */
 
 #include <cuml/metrics/metrics.hpp>
+#include <raft/distance/specializations.hpp>
 #include <raft/spatial/knn/knn.hpp>
+#include <raft/spatial/knn/specializations.hpp>
 #include <rmm/device_scalar.hpp>
 #include <rmm/device_uvector.hpp>
 #include <selection/columnWiseSort.cuh>
@@ -99,19 +101,19 @@ void run_knn(const raft::handle_t& h,
   ptrs[0]  = input;
   sizes[0] = n;
 
-  raft::spatial::knn::brute_force_knn(h,
-                                      ptrs,
-                                      sizes,
-                                      d,
-                                      input,
-                                      n,
-                                      indices,
-                                      distances,
-                                      n_neighbors,
-                                      true,
-                                      true,
-                                      nullptr,
-                                      distance_type);
+  raft::spatial::knn::brute_force_knn<int64_t, float, int>(h,
+                                                           ptrs,
+                                                           sizes,
+                                                           d,
+                                                           input,
+                                                           n,
+                                                           indices,
+                                                           distances,
+                                                           n_neighbors,
+                                                           true,
+                                                           true,
+                                                           nullptr,
+                                                           distance_type);
 }
 
 /**