Merge pull request rapidsai#140 from rapidsai/branch-24.06

Forward-merge branch-24.06 into branch-24.08

.pre-commit-config.yaml

@@ -114,7 +114,6 @@ repos:
                   setup[.]cfg$
             exclude: |
               (?x)
-                  cpp/include/cuvs/neighbors/detail/faiss_select/|
                   docs/source/sphinxext/github_link\.py|
                   cpp/cmake/modules/FindAVX\.cmake|
 

cpp/CMakeLists.txt

@@ -447,14 +447,29 @@ target_include_directories(
               "$<INSTALL_INTERFACE:include>"
 )
 
+rapids_find_package(
+  OpenMP REQUIRED
+  BUILD_EXPORT_SET cuvs-exports
+  INSTALL_EXPORT_SET cuvs-exports
+)
+
 if(NOT BUILD_CPU_ONLY)
+
+  set(CUVS_CUSOLVER_DEPENDENCY CUDA::cusolver${_ctk_static_suffix})
+  set(CUVS_CUBLAS_DEPENDENCY CUDA::cublas${_ctk_static_suffix})
+  set(CUVS_CURAND_DEPENDENCY CUDA::curand${_ctk_static_suffix})
+  set(CUVS_CUSPARSE_DEPENDENCY CUDA::cusparse${_ctk_static_suffix})
+
+  set(CUVS_CTK_MATH_DEPENDENCIES ${CUVS_CUBLAS_DEPENDENCY} ${CUVS_CUSOLVER_DEPENDENCY}
+                                 ${CUVS_CUSPARSE_DEPENDENCY} ${CUVS_CURAND_DEPENDENCY}
+  )
+
   # Keep cuVS as lightweight as possible. Only CUDA libs and rmm should be used in global target.
   target_link_libraries(
     cuvs
-    PUBLIC rmm::rmm $<$<NOT:$<BOOL:${CUDA_STATIC_RUNTIME}>>:raft::raft>
-           $<$<NOT:$<BOOL:${CUDA_STATIC_RUNTIME}>>:raft::compiled>
-    PRIVATE $<$<BOOL:${CUDA_STATIC_RUNTIME}>:raft::raft>
-            $<$<BOOL:${CUDA_STATIC_RUNTIME}>:raft::compiled_static> nvidia::cutlass::cutlass
+    PUBLIC rmm::rmm raft::raft
+    PRIVATE nvidia::cutlass::cutlass ${CUVS_CTK_MATH_DEPENDENCIES}
+            $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX>
   )
 endif()
 
@@ -559,7 +574,7 @@ if(BUILD_C_LIBRARY)
   target_link_libraries(
     cuvs_c
     PUBLIC cuvs::cuvs
-    PRIVATE $<$<BOOL:${CUDA_STATIC_RUNTIME}>:raft::raft>
+    PRIVATE raft::raft ${CUVS_CTK_MATH_DEPENDENCIES}
   )
 
   # ensure CUDA symbols aren't relocated to the middle of the debug build binaries

cpp/cmake/thirdparty/get_raft.cmake

@@ -17,7 +17,7 @@ set(RAFT_FORK "rapidsai")
 set(RAFT_PINNED_TAG "branch-${RAPIDS_VERSION_MAJOR_MINOR}")
 
 function(find_and_configure_raft)
-    set(oneValueArgs VERSION FORK PINNED_TAG COMPILE_LIBRARY USE_RAFT_STATIC ENABLE_NVTX ENABLE_MNMG_DEPENDENCIES)
+    set(oneValueArgs VERSION FORK PINNED_TAG USE_RAFT_STATIC ENABLE_NVTX ENABLE_MNMG_DEPENDENCIES)
     cmake_parse_arguments(PKG "${options}" "${oneValueArgs}"
             "${multiValueArgs}" ${ARGN} )
 
@@ -31,16 +31,6 @@ function(find_and_configure_raft)
 
     set(RAFT_COMPONENTS "")
 
-    if(PKG_COMPILE_LIBRARY)
-      if(NOT PKG_USE_RAFT_STATIC)
-        string(APPEND RAFT_COMPONENTS " compiled")
-        set(RAFT_COMPILED_LIB raft::compiled PARENT_SCOPE)
-      else()
-        string(APPEND RAFT_COMPONENTS " compiled_static")
-        set(RAFT_COMPILED_LIB raft::compiled_static PARENT_SCOPE)
-      endif()
-    endif()
-
     if(PKG_ENABLE_MNMG_DEPENDENCIES)
         string(APPEND RAFT_COMPONENTS " distributed")
     endif()
@@ -62,7 +52,7 @@ function(find_and_configure_raft)
               "BUILD_PRIMS_BENCH OFF"
               "BUILD_ANN_BENCH OFF"
               "RAFT_NVTX ${PKG_ENABLE_NVTX}"
-              "RAFT_COMPILE_LIBRARY ${PKG_COMPILE_LIBRARY}"
+              "RAFT_COMPILE_LIBRARY OFF"
             )
 endfunction()
 
@@ -72,7 +62,6 @@ endfunction()
 find_and_configure_raft(VERSION  ${RAFT_VERSION}.00
         FORK                     ${RAFT_FORK}
         PINNED_TAG               ${RAFT_PINNED_TAG}
-        COMPILE_LIBRARY          ON
         ENABLE_MNMG_DEPENDENCIES OFF
         ENABLE_NVTX              OFF
         USE_RAFT_STATIC ${CUVS_USE_RAFT_STATIC}

cpp/include/cuvs/cluster/agglomerative.hpp

@@ -16,7 +16,7 @@
 
 #pragma once
 
-#include <cuvs/distance/distance_types.hpp>
+#include <cuvs/distance/distance.hpp>
 #include <optional>
 #include <raft/core/device_mdspan.hpp>
 #include <raft/core/resources.hpp>

cpp/include/cuvs/cluster/kmeans.hpp

@@ -14,7 +14,7 @@
  * limitations under the License.
  */
 #pragma once
-#include <cuvs/distance/distance_types.hpp>
+#include <cuvs/distance/distance.hpp>
 #include <raft/core/device_mdspan.hpp>
 #include <raft/core/host_mdspan.hpp>
 #include <raft/core/resources.hpp>

cpp/include/cuvs/distance/pairwise_distance.hpp → cpp/include/cuvs/distance/distance.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -16,13 +16,54 @@
 
 #pragma once
 
+#include "distance.h"
+
 #include <cstdint>
-#include <cuvs/distance/distance_types.hpp>
 #include <raft/core/device_mdspan.hpp>
 #include <raft/core/resources.hpp>
 
 namespace cuvs::distance {
 
+using DistanceType = cuvsDistanceType;
+
+/**
+ * Whether minimal distance corresponds to similar elements (using the given metric).
+ */
+inline bool is_min_close(DistanceType metric)
+{
+  bool select_min;
+  switch (metric) {
+    case DistanceType::InnerProduct:
+      // Similarity metrics have the opposite meaning, i.e. nearest neighbors are those with larger
+      // similarity (See the same logic at cpp/include/raft/sparse/spatial/detail/knn.cuh:362
+      // {perform_k_selection})
+      select_min = false;
+      break;
+    default: select_min = true;
+  }
+  return select_min;
+}
+
+namespace kernels {
+enum KernelType { LINEAR, POLYNOMIAL, RBF, TANH };
+
+/**
+ * Parameters for kernel matrices.
+ * The following kernels are implemented:
+ * - LINEAR \f[ K(x_1,x_2) = <x_1,x_2>, \f] where \f$< , >\f$ is the dot product
+ * - POLYNOMIAL \f[ K(x_1, x_2) = (\gamma <x_1,x_2> + \mathrm{coef0})^\mathrm{degree} \f]
+ * - RBF \f[ K(x_1, x_2) = \exp(- \gamma |x_1-x_2|^2) \f]
+ * - TANH \f[ K(x_1, x_2) = \tanh(\gamma <x_1,x_2> + \mathrm{coef0}) \f]
+ */
+struct KernelParams {
+  // Kernel function parameters
+  KernelType kernel;  //!< Type of the kernel function
+  int degree;         //!< Degree of polynomial kernel (ignored by others)
+  double gamma;       //!< multiplier in the
+  double coef0;       //!< additive constant in poly and tanh kernels
+};
+}  // end namespace kernels
+
 /**
  * @defgroup pairwise_distance Pairwise Distances API
  * @{
@@ -37,7 +78,7 @@ namespace cuvs::distance {
  * @code{.cpp}
  * #include <raft/core/resources.hpp>
  * #include <raft/core/device_mdarray.hpp>
- * #include <cuvs/distance/pairwise_distance.hpp>
+ * #include <cuvs/distance/distance.hpp>
  *
  * raft::resources handle;
  * int n_samples = 5000;
@@ -81,7 +122,7 @@ void pairwise_distance(
  * @code{.cpp}
  * #include <raft/core/resources.hpp>
  * #include <raft/core/device_mdarray.hpp>
- * #include <cuvs/distance/pairwise_distance.hpp>
+ * #include <cuvs/distance/distance.hpp>
  *
  * raft::resources handle;
  * int n_samples = 5000;
@@ -125,7 +166,7 @@ void pairwise_distance(
  * @code{.cpp}
  * #include <raft/core/resources.hpp>
  * #include <raft/core/device_mdarray.hpp>
- * #include <cuvs/distance/pairwise_distance.hpp>
+ * #include <cuvs/distance/distance.hpp>
  *
  * raft::resources handle;
  * int n_samples = 5000;
@@ -168,7 +209,7 @@ void pairwise_distance(
  * @code{.cpp}
  * #include <raft/core/resources.hpp>
  * #include <raft/core/device_mdarray.hpp>
- * #include <cuvs/distance/pairwise_distance.hpp>
+ * #include <cuvs/distance/distance.hpp>
  *
  * raft::resources handle;
  * int n_samples = 5000;
@@ -205,4 +246,4 @@ void pairwise_distance(
 
 /** @} */  // end group pairwise_distance_runtime
 
-}  // namespace cuvs::distance
+};  // namespace cuvs::distance

cpp/include/cuvs/neighbors/brute_force.h

@@ -17,7 +17,7 @@
 #pragma once
 
 #include <cuvs/core/c_api.h>
-#include <cuvs/distance/distance_types.h>
+#include <cuvs/distance/distance.h>
 #include <dlpack/dlpack.h>
 #include <stdint.h>