Merge pull request #1350 from cjnolet/branch-23.04-merge_into_branch-…

…23.06 Forward merge branch 23.04 into 23.06
rapidsai · Mar 19, 2023 · f66198b · f66198b
2 parents 1053f29 + 1a3c028
commit f66198b
Show file tree

Hide file tree

Showing 168 changed files with 8,324 additions and 2,101 deletions.
diff --git a/.gitignore b/.gitignore
@@ -24,6 +24,7 @@ log
 .DS_Store
 dask-worker-space/
 *.egg-info/
+*.bin
 
 ## scikit-build
 _skbuild

diff --git a/README.md b/README.md
@@ -46,7 +46,7 @@ RAFT relies heavily on RMM which eases the burden of configuring different alloc
 
 ### Multi-dimensional Arrays
 
-The APIs in RAFT currently accept raw pointers to device memory and we are in the process of simplifying the APIs with the [mdspan](https://arxiv.org/abs/2010.06474) multi-dimensional array view for representing data in higher dimensions similar to the `ndarray` in the Numpy Python library. RAFT also contains the corresponding owning `mdarray` structure, which simplifies the allocation and management of multi-dimensional data in both host and device (GPU) memory.
+The APIs in RAFT accept the [mdspan](https://arxiv.org/abs/2010.06474) multi-dimensional array view for representing data in higher dimensions similar to the `ndarray` in the Numpy Python library. RAFT also contains the corresponding owning `mdarray` structure, which simplifies the allocation and management of multi-dimensional data in both host and device (GPU) memory.
 
 The `mdarray` forms a convenience layer over RMM and can be constructed in RAFT using a number of different helper functions:
 

diff --git a/ci/checks/copyright.py b/ci/checks/copyright.py
@@ -37,7 +37,7 @@
     re.compile(r"setup[.]cfg$"),
     re.compile(r"meta[.]yaml$")
 ]
-ExemptFiles = ["cpp/include/raft/spatial/knn/detail/faiss_select/"]
+ExemptFiles = ["cpp/include/raft/neighbors/detail/faiss_select/"]
 
 # this will break starting at year 10000, which is probably OK :)
 CheckSimple = re.compile(

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
@@ -223,6 +223,10 @@ target_link_libraries(
 )
 
 target_compile_features(raft INTERFACE cxx_std_17 $<BUILD_INTERFACE:cuda_std_17>)
+target_compile_options(
+  raft INTERFACE $<$<COMPILE_LANG_AND_ID:CUDA,NVIDIA>:--expt-extended-lambda
+                 --expt-relaxed-constexpr>
+)
 
 # Endian detection
 include(TestBigEndian)
@@ -364,6 +368,17 @@ if(RAFT_COMPILE_DIST_LIBRARY)
     src/distance/matrix/specializations/detail/select_k_float_int64_t.cu
     src/distance/matrix/specializations/detail/select_k_half_uint32_t.cu
     src/distance/matrix/specializations/detail/select_k_half_int64_t.cu
+    src/distance/neighbors/ivf_flat_search.cu
+    src/distance/neighbors/ivf_flat_build.cu
+    src/distance/neighbors/specializations/ivfflat_build_float_int64_t.cu
+    src/distance/neighbors/specializations/ivfflat_build_int8_t_int64_t.cu
+    src/distance/neighbors/specializations/ivfflat_build_uint8_t_int64_t.cu
+    src/distance/neighbors/specializations/ivfflat_extend_float_int64_t.cu
+    src/distance/neighbors/specializations/ivfflat_extend_int8_t_int64_t.cu
+    src/distance/neighbors/specializations/ivfflat_extend_uint8_t_int64_t.cu
+    src/distance/neighbors/specializations/ivfflat_search_float_int64_t.cu
+    src/distance/neighbors/specializations/ivfflat_search_int8_t_int64_t.cu
+    src/distance/neighbors/specializations/ivfflat_search_uint8_t_int64_t.cu
     src/distance/neighbors/ivfpq_build.cu
     src/distance/neighbors/ivfpq_deserialize.cu
     src/distance/neighbors/ivfpq_serialize.cu

diff --git a/cpp/bench/cluster/kmeans.cu b/cpp/bench/cluster/kmeans.cu
@@ -38,7 +38,13 @@ inline auto operator<<(std::ostream& os, const KMeansBenchParams& p) -> std::ost
 
 template <typename T, typename IndexT = int>
 struct KMeans : public BlobsFixture<T, IndexT> {
-  KMeans(const KMeansBenchParams& p) : BlobsFixture<T, IndexT>(p.data, p.blobs), params(p) {}
+  KMeans(const KMeansBenchParams& p)
+    : BlobsFixture<T, IndexT>(p.data, p.blobs),
+      params(p),
+      centroids(this->handle),
+      labels(this->handle)
+  {
+  }
 
   void run_benchmark(::benchmark::State& state) override
   {

diff --git a/cpp/bench/cluster/kmeans_balanced.cu b/cpp/bench/cluster/kmeans_balanced.cu
@@ -32,7 +32,7 @@ struct KMeansBalancedBenchParams {
 
 template <typename T, typename IndexT = int>
 struct KMeansBalanced : public fixture {
-  KMeansBalanced(const KMeansBalancedBenchParams& p) : params(p) {}
+  KMeansBalanced(const KMeansBalancedBenchParams& p) : params(p), X(handle), centroids(handle) {}
 
   void run_benchmark(::benchmark::State& state) override
   {

diff --git a/cpp/bench/common/benchmark.hpp b/cpp/bench/common/benchmark.hpp
@@ -172,7 +172,10 @@ struct BlobsParams {
 template <typename T, typename IndexT = int>
 class BlobsFixture : public fixture {
  public:
-  BlobsFixture(const DatasetParams dp, const BlobsParams bp) : data_params(dp), blobs_params(bp) {}
+  BlobsFixture(const DatasetParams dp, const BlobsParams bp)
+    : data_params(dp), blobs_params(bp), X(this->handle)
+  {
+  }
 
   virtual void run_benchmark(::benchmark::State& state) = 0;
 

diff --git a/cpp/bench/distance/fused_l2_nn.cu b/cpp/bench/distance/fused_l2_nn.cu
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-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.
@@ -36,7 +36,16 @@ inline auto operator<<(std::ostream& os, const fusedl2nn_inputs& p) -> std::ostr
 
 template <typename DataT, typename IdxT, typename OutT>
 struct fusedl2nn : public fixture {
-  fusedl2nn(const fusedl2nn_inputs& p) : params(p) {}
+  fusedl2nn(const fusedl2nn_inputs& p)
+    : params(p),
+      workspace(this->handle),
+      x(this->handle),
+      y(this->handle),
+      x_norm(this->handle),
+      y_norm(this->handle),
+      out(this->handle)
+  {
+  }
 
   void allocate_data(const ::benchmark::State& state) override
   {

diff --git a/cpp/bench/matrix/argmin.cu b/cpp/bench/matrix/argmin.cu
@@ -30,7 +30,7 @@ struct ArgminParams {
 
 template <typename T, typename OutT, typename IdxT>
 struct Argmin : public fixture {
-  Argmin(const ArgminParams<IdxT>& p) : params(p) {}
+  Argmin(const ArgminParams<IdxT>& p) : params(p), matrix(this->handle), indices(this->handle) {}
 
   void allocate_data(const ::benchmark::State& state) override
   {

diff --git a/cpp/bench/matrix/gather.cu b/cpp/bench/matrix/gather.cu
@@ -37,7 +37,10 @@ inline auto operator<<(std::ostream& os, const GatherParams<IdxT>& p) -> std::os
 
 template <typename T, typename MapT, typename IdxT, bool Conditional = false>
 struct Gather : public fixture {
-  Gather(const GatherParams<IdxT>& p) : params(p) {}
+  Gather(const GatherParams<IdxT>& p)
+    : params(p), matrix(this->handle), map(this->handle), out(this->handle), stencil(this->handle)
+  {
+  }
 
   void allocate_data(const ::benchmark::State& state) override
   {

diff --git a/cpp/bench/neighbors/knn.cuh b/cpp/bench/neighbors/knn.cuh
@@ -178,7 +178,6 @@ struct ivf_pq_knn {
   {
     index_params.n_lists = 4096;
     index_params.metric  = raft::distance::DistanceType::L2Expanded;
-
     auto data_view = raft::make_device_matrix_view<const ValT, IdxT>(data, ps.n_samples, ps.n_dims);
     index.emplace(raft::neighbors::ivf_pq::build(handle, index_params, data_view));
   }
@@ -189,13 +188,12 @@ struct ivf_pq_knn {
               IdxT* out_idxs)
   {
     search_params.n_probes = 20;
-
     auto queries_view =
       raft::make_device_matrix_view<const ValT, IdxT>(search_items, ps.n_queries, ps.n_dims);
     auto idxs_view  = raft::make_device_matrix_view<IdxT, IdxT>(out_idxs, ps.n_queries, ps.k);
     auto dists_view = raft::make_device_matrix_view<dist_t, IdxT>(out_dists, ps.n_queries, ps.k);
     raft::neighbors::ivf_pq::search(
-      handle, search_params, *index, queries_view, ps.k, idxs_view, dists_view);
+      handle, search_params, *index, queries_view, idxs_view, dists_view);
   }
 };