Add random subsampling for IVF methods

rapidsai · rapids-bot · Jan 23, 2024 · Jan 3, 2024 · Jan 3, 2024 · Jan 18, 2024
commit a455b08b4a70b93da0c9b548e269e8c82b81f696
diff --git a/cpp/bench/ann/src/raft/raft_ann_bench_param_parser.h b/cpp/bench/ann/src/raft/raft_ann_bench_param_parser.h
@@ -55,6 +55,7 @@ void parse_build_param(const nlohmann::json& conf,
   param.n_lists = conf.at("nlist");
   if (conf.contains("niter")) { param.kmeans_n_iters = conf.at("niter"); }
   if (conf.contains("ratio")) { param.kmeans_trainset_fraction = 1.0 / (double)conf.at("ratio"); }
+  if (conf.contains("random_seed")) { param.random_seed = conf.at("random_seed"); }
 }
 
 template <typename T, typename IdxT>
@@ -87,6 +88,7 @@ void parse_build_param(const nlohmann::json& conf,
                                "', should be either 'cluster' or 'subspace'");
     }
   }
+  if (conf.contains("random_seed")) { param.random_seed = conf.at("random_seed"); }
 }
 
 template <typename T, typename IdxT>

@@ -364,25 +364,19 @@ inline auto build(raft::resources const& handle,
     auto trainset_ratio = std::max<size_t>(
       1, n_rows / std::max<size_t>(params.kmeans_trainset_fraction * n_rows, index.n_lists()));
     auto n_rows_train = n_rows / trainset_ratio;
-    rmm::device_uvector<T> trainset(n_rows_train * index.dim(), stream);
-    // TODO: a proper sampling
-    RAFT_CUDA_TRY(cudaMemcpy2DAsync(trainset.data(),
-                                    sizeof(T) * index.dim(),
-                                    dataset,
-                                    sizeof(T) * index.dim() * trainset_ratio,
-                                    sizeof(T) * index.dim(),
-                                    n_rows_train,
-                                    cudaMemcpyDefault,
-                                    stream));
-    auto trainset_const_view =
-      raft::make_device_matrix_view<const T, IdxT>(trainset.data(), n_rows_train, index.dim());
+    auto trainset     = make_device_matrix<T, IdxT>(handle, n_rows_train, index.dim());
+    raft::spatial::knn::detail::utils::subsample(
+      handle, dataset, n_rows, trainset.view(), params.random_seed);
     auto centers_view = raft::make_device_matrix_view<float, IdxT>(
       index.centers().data_handle(), index.n_lists(), index.dim());
     raft::cluster::kmeans_balanced_params kmeans_params;
     kmeans_params.n_iters = params.kmeans_n_iters;
     kmeans_params.metric  = index.metric();
-    raft::cluster::kmeans_balanced::fit(
-      handle, kmeans_params, trainset_const_view, centers_view, utils::mapping<float>{});
+    raft::cluster::kmeans_balanced::fit(handle,
+                                        kmeans_params,
+                                        make_const_mdspan(trainset.view()),
+                                        centers_view,
+                                        utils::mapping<float>{});
   }
 
   // add the data if necessary

@@ -25,9 +25,11 @@
 
 #include <raft/cluster/kmeans_balanced.cuh>
 #include <raft/core/device_mdarray.hpp>
+#include <raft/core/device_resources.hpp>
 #include <raft/core/logger.hpp>
 #include <raft/core/nvtx.hpp>
 #include <raft/core/operators.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/distance/distance_types.hpp>
 #include <raft/linalg/add.cuh>
@@ -46,7 +48,6 @@
 #include <raft/util/pow2_utils.cuh>
 #include <raft/util/vectorized.cuh>
 
-#include <raft/core/resource/device_memory_resource.hpp>
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/device_uvector.hpp>
 #include <rmm/mr/device/managed_memory_resource.hpp>
@@ -1759,71 +1760,41 @@ auto build(raft::resources const& handle,
       size_t(n_rows) / std::max<size_t>(params.kmeans_trainset_fraction * n_rows, index.n_lists()));
     size_t n_rows_train = n_rows / trainset_ratio;
 
-    auto* device_memory = resource::get_workspace_resource(handle);
-    rmm::mr::managed_memory_resource managed_memory_upstream;
+    auto* device_mr = resource::get_workspace_resource(handle);
+    rmm::mr::managed_memory_resource managed_mr;
 
     // Besides just sampling, we transform the input dataset into floats to make it easier
     // to use gemm operations from cublas.
-    rmm::device_uvector<float> trainset(n_rows_train * index.dim(), stream, device_memory);
-    // TODO: a proper sampling
+    auto trainset =
+      make_device_mdarray<float>(handle, device_mr, make_extents<IdxT>(n_rows_train, dim));
+
     if constexpr (std::is_same_v<T, float>) {
-      RAFT_CUDA_TRY(cudaMemcpy2DAsync(trainset.data(),
-                                      sizeof(T) * index.dim(),
-                                      dataset,
-                                      sizeof(T) * index.dim() * trainset_ratio,
-                                      sizeof(T) * index.dim(),
-                                      n_rows_train,
-                                      cudaMemcpyDefault,
-                                      stream));
+      raft::spatial::knn::detail::utils::subsample(
+        handle, dataset, n_rows, trainset.view(), params.random_seed);
     } else {
-      size_t dim = index.dim();
-      cudaPointerAttributes dataset_attr;
-      RAFT_CUDA_TRY(cudaPointerGetAttributes(&dataset_attr, dataset));
-      if (dataset_attr.devicePointer != nullptr) {
-        // data is available on device: just run the kernel to copy and map the data
-        auto p = reinterpret_cast<T*>(dataset_attr.devicePointer);
-        auto trainset_view =
-          raft::make_device_vector_view<float, IdxT>(trainset.data(), dim * n_rows_train);
-        linalg::map_offset(handle, trainset_view, [p, trainset_ratio, dim] __device__(size_t i) {
-          auto col = i % dim;
-          return utils::mapping<float>{}(p[(i - col) * size_t(trainset_ratio) + col]);
-        });
-      } else {
-        // data is not available: first copy, then map inplace
-        auto trainset_tmp = reinterpret_cast<T*>(reinterpret_cast<uint8_t*>(trainset.data()) +
-                                                 (sizeof(float) - sizeof(T)) * index.dim());
-        // We copy the data in strides, one row at a time, and place the smaller rows of type T
-        // at the end of float rows.
-        RAFT_CUDA_TRY(cudaMemcpy2DAsync(trainset_tmp,
-                                        sizeof(float) * index.dim(),
-                                        dataset,
-                                        sizeof(T) * index.dim() * trainset_ratio,
-                                        sizeof(T) * index.dim(),
-                                        n_rows_train,
-                                        cudaMemcpyDefault,
-                                        stream));
-        // Transform the input `{T -> float}`, one row per warp.
-        // The threads in each warp copy the data synchronously; this and the layout of the data
-        // (content is aligned to the end of the rows) together allow doing the transform in-place.
-        copy_warped(trainset.data(),
-                    index.dim(),
-                    trainset_tmp,
-                    index.dim() * sizeof(float) / sizeof(T),
-                    index.dim(),
-                    n_rows_train,
-                    stream);
-      }
+      // TODO(tfeher): Enable codebook generation with any type T, and then remove
+      // trainset tmp.
+      auto trainset_tmp =
+        make_device_mdarray<T>(handle, device_mr, make_extents<IdxT>(n_rows_train, dim));
+      raft::spatial::knn::detail::utils::subsample(
+        handle, dataset, n_rows, trainset_tmp.view(), params.random_seed);
+      cudaDeviceSynchronize();
+      RAFT_LOG_INFO("Subsampling done, converting to float");
+      raft::linalg::unaryOp(trainset.data_handle(),
+                            trainset_tmp.data_handle(),
+                            trainset.size(),
+                            utils::mapping<float>{},  // raft::cast_op<float>(),
+                            raft::resource::get_cuda_stream(handle));
     }
 
     // NB: here cluster_centers is used as if it is [n_clusters, data_dim] not [n_clusters,
     // dim_ext]!
     rmm::device_uvector<float> cluster_centers_buf(
-      index.n_lists() * index.dim(), stream, device_memory);
+      index.n_lists() * index.dim(), stream, device_mr);
     auto cluster_centers = cluster_centers_buf.data();
 
     // Train balanced hierarchical kmeans clustering
-    auto trainset_const_view =
-      raft::make_device_matrix_view<const float, IdxT>(trainset.data(), n_rows_train, index.dim());
+    auto trainset_const_view = raft::make_const_mdspan(trainset.view());
     auto centers_view =
       raft::make_device_matrix_view<float, IdxT>(cluster_centers, index.n_lists(), index.dim());
     raft::cluster::kmeans_balanced_params kmeans_params;
@@ -1833,7 +1804,7 @@ auto build(raft::resources const& handle,
       handle, kmeans_params, trainset_const_view, centers_view, utils::mapping<float>{});
 
     // Trainset labels are needed for training PQ codebooks
-    rmm::device_uvector<uint32_t> labels(n_rows_train, stream, device_memory);
+    rmm::device_uvector<uint32_t> labels(n_rows_train, stream, device_mr);
     auto centers_const_view = raft::make_device_matrix_view<const float, IdxT>(
       cluster_centers, index.n_lists(), index.dim());
     auto labels_view = raft::make_device_vector_view<uint32_t, IdxT>(labels.data(), n_rows_train);
@@ -1859,19 +1830,19 @@ auto build(raft::resources const& handle,
         train_per_subset(handle,
                          index,
                          n_rows_train,
-                         trainset.data(),
+                         trainset.data_handle(),
                          labels.data(),
                          params.kmeans_n_iters,
-                         &managed_memory_upstream);
+                         &managed_mr);
         break;
       case codebook_gen::PER_CLUSTER:
         train_per_cluster(handle,
                           index,
                           n_rows_train,
-                          trainset.data(),
+                          trainset.data_handle(),
                           labels.data(),
                           params.kmeans_n_iters,
-                          &managed_memory_upstream);
+                          &managed_mr);
         break;
       default: RAFT_FAIL("Unreachable code");
     }

diff --git a/cpp/include/raft/neighbors/ivf_flat_types.hpp b/cpp/include/raft/neighbors/ivf_flat_types.hpp
@@ -76,6 +76,12 @@ struct index_params : ann::index_params {
    * flag to `true` if you prefer to use as little GPU memory for the database as possible.
    */
   bool conservative_memory_allocation = false;
+  /**
+   * Seed used for random sampling if kmeans_trainset_fraction < 1.
+   *
+   * Value -1 disables random sampling, and results in sampling with a fixed stride.
+   */
+  int random_seed = 0;
 };
 
 struct search_params : ann::search_params {

diff --git a/cpp/include/raft/neighbors/ivf_pq_types.hpp b/cpp/include/raft/neighbors/ivf_pq_types.hpp
@@ -105,6 +105,13 @@ struct index_params : ann::index_params {
    * flag to `true` if you prefer to use as little GPU memory for the database as possible.
    */
   bool conservative_memory_allocation = false;
+
+  /**
+   * Seed used for random sampling if kmeans_trainset_fraction < 1.
+   *
+   * Value -1 disables random sampling, and results in sampling with a fixed stride.
+   */
+  int random_seed = 0;
 };
 
 struct search_params : ann::search_params {

@@ -16,8 +16,14 @@
 
 #pragma once
 
+#include <raft/core/device_mdarray.hpp>
+#include <raft/core/device_mdspan.hpp>
+#include <raft/core/device_resources.hpp>
+#include <raft/core/host_mdarray.hpp>
 #include <raft/core/logger.hpp>
 #include <raft/distance/distance_types.hpp>
+#include <raft/matrix/copy.cuh>
+#include <raft/random/sample_without_replacement.cuh>
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/cudart_utils.hpp>
 #include <raft/util/integer_utils.hpp>
@@ -30,6 +36,10 @@
 #include <optional>
 
 #include <cuda_fp16.hpp>
+#include <thrust/copy.h>
+#include <thrust/device_ptr.h>
+#include <thrust/iterator/counting_iterator.h>
+#include <thrust/sort.h>
 
 namespace raft::spatial::knn::detail::utils {
 
@@ -573,4 +583,85 @@ struct batch_load_iterator {
   size_type cur_pos_;
 };
 
+template <typename IdxT>
+auto get_subsample_indices(raft::resources const& res, IdxT n_samples, IdxT n_subsamples, int seed)
+  -> raft::device_vector<IdxT, IdxT>
+{
+  RAFT_EXPECTS(n_subsamples <= n_samples, "Cannot have more training samples than dataset vectors");
+
+  auto data_indices = raft::make_device_vector<IdxT, IdxT>(res, n_samples);
+  thrust::counting_iterator<IdxT> first(0);
+  thrust::device_ptr<IdxT> ptr(data_indices.data_handle());
+  thrust::copy(raft::resource::get_thrust_policy(res), first, first + n_samples, ptr);
+  raft::random::RngState rng(seed);
+  auto train_indices = raft::make_device_vector<IdxT, IdxT>(res, n_subsamples);
+  raft::random::sample_without_replacement(res,
+                                           rng,
+                                           raft::make_const_mdspan(data_indices.view()),
+                                           std::nullopt,
+                                           train_indices.view(),
+                                           std::nullopt);
+
+  thrust::sort(resource::get_thrust_policy(res),
+               train_indices.data_handle(),
+               train_indices.data_handle() + n_subsamples);
+  return train_indices;
+}
+
+/** Subsample the dataset to create a training set*/
+template <typename T, typename IdxT = int64_t>
+void subsample(raft::resources const& res,
+               const T* input,
+               IdxT n_samples,
+               raft::device_matrix_view<T, IdxT> output,
+               int seed)
+{
+  int64_t n_dim   = output.extent(1);
+  int64_t n_train = output.extent(0);
+  if (seed == -1 || n_train == n_samples) {
+    IdxT trainset_ratio = n_samples / n_train;
+    RAFT_LOG_INFO("Fixed stride subsampling");
+    RAFT_CUDA_TRY(cudaMemcpy2DAsync(output.data_handle(),
+                                    sizeof(T) * n_dim,
+                                    input,
+                                    sizeof(T) * n_dim * trainset_ratio,
+                                    sizeof(T) * n_dim,
+                                    n_train,
+                                    cudaMemcpyDefault,
+                                    resource::get_cuda_stream(res)));
+    return;
+  }
+  RAFT_LOG_DEBUG("Random subsampling");
+  raft::device_vector<IdxT, IdxT> train_indices =
+    get_subsample_indices<IdxT>(res, n_samples, n_train, seed);
+
+  cudaPointerAttributes attr;
+  RAFT_CUDA_TRY(cudaPointerGetAttributes(&attr, input));
+  T* ptr = reinterpret_cast<T*>(attr.devicePointer);
+  if (ptr != nullptr) {
+    raft::matrix::copy_rows(res,
+                            raft::make_device_matrix_view<const T, IdxT>(ptr, n_samples, n_dim),
+                            output,
+                            raft::make_const_mdspan(train_indices.view()));
+  } else {
+    auto dataset            = raft::make_host_matrix_view<const T, IdxT>(input, n_samples, n_dim);
+    auto train_indices_host = raft::make_host_vector<IdxT, IdxT>(n_train);
+    raft::copy(train_indices_host.data_handle(),
+               train_indices.data_handle(),
+               n_train,
+               resource::get_cuda_stream(res));
+    resource::sync_stream(res);
+    auto out_tmp = raft::make_host_matrix<T, IdxT>(n_train, n_dim);
+#pragma omp parallel for
+    for (IdxT i = 0; i < n_train; i++) {
+      IdxT in_idx = train_indices_host(i);
+      for (IdxT k = 0; k < n_dim; k++) {
+        out_tmp(i, k) = dataset(in_idx, k);
+      }
+    }
+    raft::copy(
+      output.data_handle(), out_tmp.data_handle(), output.size(), resource::get_cuda_stream(res));
+    resource::sync_stream(res);
+  }
+}
 }  // namespace raft::spatial::knn::detail::utils
diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh
@@ -73,14 +73,15 @@ struct AnnIvfFlatInputs {
   raft::distance::DistanceType metric;
   bool adaptive_centers;
   bool host_dataset;
+  int seed;
 };
 
 template <typename IdxT>
 ::std::ostream& operator<<(::std::ostream& os, const AnnIvfFlatInputs<IdxT>& p)
 {
   os << "{ " << p.num_queries << ", " << p.num_db_vecs << ", " << p.dim << ", " << p.k << ", "
      << p.nprobe << ", " << p.nlist << ", " << static_cast<int>(p.metric) << ", "
-     << p.adaptive_centers << ", " << p.host_dataset << '}' << std::endl;
+     << p.adaptive_centers << ", " << p.host_dataset << "," << p.seed << '}' << std::endl;
   return os;
 }
 
@@ -178,6 +179,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
         index_params.add_data_on_build        = false;
         index_params.kmeans_trainset_fraction = 0.5;
         index_params.metric_arg               = 0;
+        index_params.random_seed              = ps.seed;
 
         ivf_flat::index<DataT, IdxT> idx(handle_, index_params, ps.dim);
         ivf_flat::index<DataT, IdxT> index_2(handle_, index_params, ps.dim);
@@ -327,6 +329,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
     index_params.add_data_on_build        = false;
     index_params.kmeans_trainset_fraction = 1.0;
     index_params.metric_arg               = 0;
+    index_params.random_seed              = ps.seed;
 
     auto database_view = raft::make_device_matrix_view<const DataT, IdxT>(
       (const DataT*)database.data(), ps.num_db_vecs, ps.dim);
@@ -497,6 +500,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
         index_params.add_data_on_build        = true;
         index_params.kmeans_trainset_fraction = 0.5;
         index_params.metric_arg               = 0;
+        index_params.random_seed              = ps.seed;
 
         // Create IVF Flat index
         auto database_view = raft::make_device_matrix_view<const DataT, IdxT>(
@@ -607,6 +611,8 @@ const std::vector<AnnIvfFlatInputs<int64_t>> inputs = {
   {20, 100000, 16, 10, 20, 1024, raft::distance::DistanceType::L2Expanded, true},
   {1000, 100000, 16, 10, 20, 1024, raft::distance::DistanceType::L2Expanded, true},
   {10000, 131072, 8, 10, 20, 1024, raft::distance::DistanceType::L2Expanded, false},
+  {10000, 1000000, 96, 10, 20, 1024, raft::distance::DistanceType::L2Expanded, false, true, -1},
+  {10000, 1000000, 96, 10, 20, 1024, raft::distance::DistanceType::L2Expanded, false, false, -1},
 
   // host input data
   {1000, 10000, 16, 10, 40, 1024, raft::distance::DistanceType::L2Expanded, false, true},

@@ -68,7 +68,7 @@ struct ivf_pq_inputs {
   ivf_pq_inputs()
   {
     index_params.n_lists                  = max(32u, min(1024u, num_db_vecs / 128u));
-    index_params.kmeans_trainset_fraction = 1.0;
+    index_params.kmeans_trainset_fraction = 0.95;
   }
 };