rapidsai · rapids-bot · Aug 24, 2022 · May 13, 2022 · May 16, 2022 · May 16, 2022
diff --git a/cpp/bench/common/benchmark.hpp b/cpp/bench/common/benchmark.hpp
@@ -40,7 +40,7 @@ struct using_pool_memory_res {
  private:
   rmm::mr::device_memory_resource* orig_res_;
   rmm::mr::cuda_memory_resource cuda_res_;
-  rmm::mr::pool_memory_resource<rmm::mr::cuda_memory_resource> pool_res_;
+  rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_res_;
 
  public:
   using_pool_memory_res(size_t initial_size, size_t max_size)

diff --git a/cpp/bench/spatial/knn.cu b/cpp/bench/spatial/knn.cu
@@ -14,25 +14,22 @@
  * limitations under the License.
  */
 
-#include <optional>
-
 #include <common/benchmark.hpp>
-#include <raft/spatial/knn/ann.cuh>
-#include <raft/spatial/knn/knn.cuh>
 
+#include <raft/random/rng.cuh>
+#include <raft/spatial/knn/knn.cuh>
 #if defined RAFT_NN_COMPILED
-#include <raft/spatial/knn/specializations.hpp>
+#include <raft/spatial/knn/specializations.cuh>
 #endif
 
-#include <raft/random/rng.cuh>
-#include <raft/sparse/detail/utils.h>
-
 #include <rmm/mr/device/managed_memory_resource.hpp>
 #include <rmm/mr/device/per_device_resource.hpp>
 
 #include <rmm/mr/host/new_delete_resource.hpp>
 #include <rmm/mr/host/pinned_memory_resource.hpp>
 
+#include <optional>
+
 namespace raft::bench::spatial {
 
 struct params {
@@ -41,14 +38,14 @@ struct params {
   /** Number of dimensions in the dataset. */
   size_t n_dims;
   /** The batch size -- number of KNN searches. */
-  size_t n_probes;
+  size_t n_queries;
   /** Number of nearest neighbours to find for every probe. */
   size_t k;
 };
 
 auto operator<<(std::ostream& os, const params& p) -> std::ostream&
 {
-  os << p.n_samples << "#" << p.n_dims << "#" << p.n_probes << "#" << p.k;
+  os << p.n_samples << "#" << p.n_dims << "#" << p.n_queries << "#" << p.k;
   return os;
 }
 
@@ -129,44 +126,6 @@ struct host_uvector {
   T* arr_;
 };
 
-template <typename ValT, typename IdxT>
-struct ivf_flat_knn {
-  using dist_t = float;
-
-  raft::spatial::knn::knnIndex index;
-  raft::spatial::knn::ivf_flat::index_params index_params;
-  raft::spatial::knn::ivf_flat::search_params search_params;
-  params ps;
-
-  ivf_flat_knn(const raft::handle_t& handle, const params& ps, const ValT* data) : ps(ps)
-  {
-    index_params.n_lists = 4096;
-    index_params.metric  = raft::distance::DistanceType::L2Expanded;
-    raft::spatial::knn::approx_knn_build_index<ValT, IdxT>(const_cast<raft::handle_t&>(handle),
-                                                           &index,
-                                                           index_params,
-                                                           const_cast<ValT*>(data),
-                                                           (IdxT)ps.n_samples,
-                                                           (IdxT)ps.n_dims);
-  }
-
-  void search(const raft::handle_t& handle,
-              const ValT* search_items,
-              dist_t* out_dists,
-              IdxT* out_idxs)
-  {
-    search_params.n_probes = 20;
-    raft::spatial::knn::approx_knn_search<ValT, IdxT>(const_cast<raft::handle_t&>(handle),
-                                                      out_dists,
-                                                      out_idxs,
-                                                      &index,
-                                                      search_params,
-                                                      (IdxT)ps.k,
-                                                      const_cast<ValT*>(search_items),
-                                                      (IdxT)ps.n_probes);
-  }
-};
-
 template <typename ValT, typename IdxT>
 struct brute_force_knn {
   using dist_t = ValT;
@@ -191,7 +150,7 @@ struct brute_force_knn {
                                                             sizes,
                                                             ps.n_dims,
                                                             const_cast<ValT*>(search_items),
-                                                            ps.n_probes,
+                                                            ps.n_queries,
                                                             out_idxs,
                                                             out_dists,
                                                             ps.k);
@@ -206,9 +165,9 @@ struct knn : public fixture {
       scope_(scope),
       dev_mem_res_(strategy == TransferStrategy::MANAGED),
       data_host_(0),
-      search_items_(p.n_probes * p.n_dims, stream),
-      out_dists_(p.n_probes * p.k, stream),
-      out_idxs_(p.n_probes * p.k, stream)
+      search_items_(p.n_queries * p.n_dims, stream),
+      out_dists_(p.n_queries * p.k, stream),
+      out_idxs_(p.n_queries * p.k, stream)
   {
     raft::random::RngState state{42};
     gen_data(state, search_items_, search_items_.size(), stream);
@@ -233,9 +192,8 @@ struct knn : public fixture {
                 size_t n,
                 rmm::cuda_stream_view stream)
   {
-    constexpr T kRangeMax = T(std::min<double>(
-      raft::spatial::knn::detail::utils::config<T>::kDivisor, std::numeric_limits<T>::max()));
-    constexpr T kRangeMin = std::is_signed_v<T> ? -kRangeMax : T(0);
+    constexpr T kRangeMax = std::is_integral_v<T> ? std::numeric_limits<T>::max() : T(1);
+    constexpr T kRangeMin = std::is_integral_v<T> ? std::numeric_limits<T>::min() : T(-1);
     if constexpr (std::is_integral_v<T>) {
       raft::random::uniformInt(state, vec.data(), n, kRangeMin, kRangeMax, stream);
     } else {
@@ -352,15 +310,12 @@ struct knn : public fixture {
 const std::vector<params> kInputs{
   {2000000, 128, 1000, 32}, {10000000, 128, 1000, 32}, {10000, 8192, 1000, 32}};
 
-const std::vector<TransferStrategy> kAllStrategies{TransferStrategy::NO_COPY,
-                                                   TransferStrategy::COPY_PLAIN,
-                                                   TransferStrategy::COPY_PINNED,
-                                                   TransferStrategy::MAP_PINNED,
-                                                   TransferStrategy::MANAGED};
+const std::vector<TransferStrategy> kAllStrategies{
+  TransferStrategy::NO_COPY, TransferStrategy::MAP_PINNED, TransferStrategy::MANAGED};
 const std::vector<TransferStrategy> kNoCopyOnly{TransferStrategy::NO_COPY};
 
 const std::vector<Scope> kScopeFull{Scope::BUILD_SEARCH};
-const std::vector<Scope> kAllScopes{Scope::BUILD, Scope::SEARCH, Scope::BUILD_SEARCH};
+const std::vector<Scope> kAllScopes{Scope::BUILD_SEARCH, Scope::SEARCH, Scope::BUILD};
 
 #define KNN_REGISTER(ValT, IdxT, ImplT, inputs, strats, scope)                   \
   namespace BENCHMARK_PRIVATE_NAME(knn)                                          \
@@ -370,8 +325,7 @@ const std::vector<Scope> kAllScopes{Scope::BUILD, Scope::SEARCH, Scope::BUILD_SE
   }
 
 KNN_REGISTER(float, int64_t, brute_force_knn, kInputs, kAllStrategies, kScopeFull);
-KNN_REGISTER(float, int64_t, ivf_flat_knn, kInputs, kNoCopyOnly, kAllScopes);
-KNN_REGISTER(int8_t, int64_t, ivf_flat_knn, kInputs, kNoCopyOnly, kAllScopes);
-KNN_REGISTER(uint8_t, int64_t, ivf_flat_knn, kInputs, kNoCopyOnly, kAllScopes);
+
+KNN_REGISTER(float, uint32_t, brute_force_knn, kInputs, kNoCopyOnly, kScopeFull);
 
 }  // namespace raft::bench::spatial
diff --git a/cpp/include/raft/spatial/knn/detail/haversine_distance.cuh b/cpp/include/raft/spatial/knn/detail/haversine_distance.cuh
@@ -72,7 +72,11 @@ __global__ void haversine_knn_kernel(value_idx* out_inds,
 
   faiss::gpu::
     BlockSelect<value_t, value_idx, false, faiss::gpu::Comparator<value_t>, warp_q, thread_q, tpb>
-      heap(faiss::gpu::Limits<value_t>::getMax(), -1, smemK, smemV, k);
+      heap(faiss::gpu::Limits<value_t>::getMax(),
+           std::numeric_limits<value_idx>::max(),
+           smemK,
+           smemV,
+           k);
 
   // Grid is exactly sized to rows available
   int limit = faiss::gpu::utils::roundDown(n_index_rows, faiss::gpu::kWarpSize);

diff --git a/cpp/include/raft/spatial/knn/detail/knn_brute_force_faiss.cuh b/cpp/include/raft/spatial/knn/detail/knn_brute_force_faiss.cuh
@@ -268,11 +268,11 @@ void brute_force_knn_impl(
   int device;
   RAFT_CUDA_TRY(cudaGetDevice(&device));
 
-  rmm::device_uvector<std::int64_t> trans(id_ranges->size(), userStream);
+  rmm::device_uvector<IdxType> trans(id_ranges->size(), userStream);
   raft::update_device(trans.data(), id_ranges->data(), id_ranges->size(), userStream);
 
   rmm::device_uvector<value_t> all_D(0, userStream);
-  rmm::device_uvector<std::int64_t> all_I(0, userStream);
+  rmm::device_uvector<IdxType> all_I(0, userStream);
 
   value_t* out_D = res_D;
   IdxType* out_I = res_I;
@@ -342,6 +342,8 @@ void brute_force_knn_impl(
           args.numQueries      = n;
           args.outDistances    = out_d_ptr;
           args.outIndices      = out_i_ptr;
+          args.outIndicesType  = sizeof(IdxType) == 4 ? faiss::gpu::IndicesDataType::I32
+                                                      : faiss::gpu::IndicesDataType::I64;
 
           /**
            * @todo: Until FAISS supports pluggable allocation strategies,

diff --git a/cpp/include/raft/spatial/knn/specializations/knn.cuh b/cpp/include/raft/spatial/knn/specializations/knn.cuh
@@ -50,6 +50,38 @@ extern template void brute_force_knn<long, float, unsigned int>(raft::handle_t c
                                                                 std::vector<long>* translations,
                                                                 distance::DistanceType metric,
                                                                 float metric_arg);
+
+extern template void brute_force_knn<uint32_t, float, int>(raft::handle_t const& handle,
+                                                           std::vector<float*>& input,
+                                                           std::vector<int>& sizes,
+                                                           int D,
+                                                           float* search_items,
+                                                           int n,
+                                                           uint32_t* res_I,
+                                                           float* res_D,
+                                                           int k,
+                                                           bool rowMajorIndex,
+                                                           bool rowMajorQuery,
+                                                           std::vector<uint32_t>* translations,
+                                                           distance::DistanceType metric,
+                                                           float metric_arg);
+
+extern template void brute_force_knn<uint32_t, float, unsigned int>(
+  raft::handle_t const& handle,
+  std::vector<float*>& input,
+  std::vector<unsigned int>& sizes,
+  unsigned int D,
+  float* search_items,
+  unsigned int n,
+  uint32_t* res_I,
+  float* res_D,
+  unsigned int k,
+  bool rowMajorIndex,
+  bool rowMajorQuery,
+  std::vector<uint32_t>* translations,
+  distance::DistanceType metric,
+  float metric_arg);
+
 };  // namespace knn
 };  // namespace spatial
 };  // namespace raft
diff --git a/cpp/src/nn/specializations/knn.cu b/cpp/src/nn/specializations/knn.cu
@@ -51,6 +51,36 @@ template void brute_force_knn<long, float, unsigned int>(raft::handle_t const& h
                                                          distance::DistanceType metric,
                                                          float metric_arg);
 
+template void brute_force_knn<uint32_t, float, int>(raft::handle_t const& handle,
+                                                    std::vector<float*>& input,
+                                                    std::vector<int>& sizes,
+                                                    int D,
+                                                    float* search_items,
+                                                    int n,
+                                                    uint32_t* res_I,
+                                                    float* res_D,
+                                                    int k,
+                                                    bool rowMajorIndex,
+                                                    bool rowMajorQuery,
+                                                    std::vector<uint32_t>* translations,
+                                                    distance::DistanceType metric,
+                                                    float metric_arg);
+
+template void brute_force_knn<uint32_t, float, unsigned int>(raft::handle_t const& handle,
+                                                             std::vector<float*>& input,
+                                                             std::vector<unsigned int>& sizes,
+                                                             unsigned int D,
+                                                             float* search_items,
+                                                             unsigned int n,
+                                                             uint32_t* res_I,
+                                                             float* res_D,
+                                                             unsigned int k,
+                                                             bool rowMajorIndex,
+                                                             bool rowMajorQuery,
+                                                             std::vector<uint32_t>* translations,
+                                                             distance::DistanceType metric,
+                                                             float metric_arg);
+
 };  // namespace knn
 };  // namespace spatial
 };  // namespace raft
diff --git a/cpp/test/spatial/knn.cu b/cpp/test/spatial/knn.cu
@@ -16,8 +16,8 @@
 
 #include "../test_utils.h"
 
+#include <raft/core/logger.hpp>
 #include <raft/distance/distance_type.hpp>
-
 #include <raft/spatial/knn/knn.cuh>
 #if defined RAFT_NN_COMPILED
 #include <raft/spatial/knn/specializations.cuh>
@@ -40,8 +40,9 @@ struct KNNInputs {
   std::vector<int> labels;
 };
 
+template <typename IdxT>
 __global__ void build_actual_output(
-  int* output, int n_rows, int k, const int* idx_labels, const int64_t* indices)
+  int* output, int n_rows, int k, const int* idx_labels, const IdxT* indices)
 {
   int element = threadIdx.x + blockDim.x * blockIdx.x;
   if (element >= n_rows * k) return;
@@ -60,7 +61,7 @@ __global__ void build_expected_output(int* output, int n_rows, int k, const int*
   }
 }
 
-template <typename T>
+template <typename T, typename IdxT>
 class KNNTest : public ::testing::TestWithParam<KNNInputs> {
  public:
   KNNTest()
@@ -79,9 +80,11 @@ class KNNTest : public ::testing::TestWithParam<KNNInputs> {
  protected:
   void testBruteForce()
   {
+#if (RAFT_ACTIVE_LEVEL >= RAFT_LEVEL_DEBUG)
     raft::print_device_vector("Input array: ", input_.data(), rows_ * cols_, std::cout);
-    std::cout << "K: " << k_ << "\n";
+    std::cout << "K: " << k_ << std::endl;
     raft::print_device_vector("Labels array: ", search_labels_.data(), rows_, std::cout);
+#endif
 
     std::vector<float*> input_vec;
     std::vector<int> sizes_vec;
@@ -131,7 +134,7 @@ class KNNTest : public ::testing::TestWithParam<KNNInputs> {
     RAFT_CUDA_TRY(cudaMemsetAsync(input_.data(), 0, input_.size() * sizeof(float), stream));
     RAFT_CUDA_TRY(
       cudaMemsetAsync(search_data_.data(), 0, search_data_.size() * sizeof(float), stream));
-    RAFT_CUDA_TRY(cudaMemsetAsync(indices_.data(), 0, indices_.size() * sizeof(int64_t), stream));
+    RAFT_CUDA_TRY(cudaMemsetAsync(indices_.data(), 0, indices_.size() * sizeof(IdxT), stream));
     RAFT_CUDA_TRY(cudaMemsetAsync(distances_.data(), 0, distances_.size() * sizeof(float), stream));
     RAFT_CUDA_TRY(
       cudaMemsetAsync(search_labels_.data(), 0, search_labels_.size() * sizeof(int), stream));
@@ -165,7 +168,7 @@ class KNNTest : public ::testing::TestWithParam<KNNInputs> {
   int cols_;
   rmm::device_uvector<float> input_;
   rmm::device_uvector<float> search_data_;
-  rmm::device_uvector<int64_t> indices_;
+  rmm::device_uvector<IdxT> indices_;
   rmm::device_uvector<float> distances_;
   int k_;
 
@@ -191,10 +194,13 @@ const std::vector<KNNInputs> inputs = {
    2,
    {0, 0, 0, 0, 0, 1, 1, 1, 1, 1}}};
 
-typedef KNNTest<float> KNNTestF;
-TEST_P(KNNTestF, BruteForce) { this->testBruteForce(); }
+typedef KNNTest<float, int64_t> KNNTestFint64_t;
+TEST_P(KNNTestFint64_t, BruteForce) { this->testBruteForce(); }
+typedef KNNTest<float, uint32_t> KNNTestFuint32_t;
+TEST_P(KNNTestFuint32_t, BruteForce) { this->testBruteForce(); }
 
-INSTANTIATE_TEST_CASE_P(KNNTest, KNNTestF, ::testing::ValuesIn(inputs));
+INSTANTIATE_TEST_CASE_P(KNNTest, KNNTestFint64_t, ::testing::ValuesIn(inputs));
+INSTANTIATE_TEST_CASE_P(KNNTest, KNNTestFuint32_t, ::testing::ValuesIn(inputs));
 
 }  // namespace knn
 }  // namespace spatial