Merge branch 'branch-24.02' into imp-2402-update_cagra_build_constraint

Author: cjnolet

conda/recipes/libraft/conda_build_config.yaml

@@ -71,3 +71,9 @@ cuda11_cuda_profiler_api_host_version:
 
 cuda11_cuda_profiler_api_run_version:
   - ">=11.4.240,<12"
+
+spdlog_version:
+  - ">=1.11.0,<1.12"
+
+fmt_version:
+  - ">=9.1.0,<10"

conda/recipes/libraft/meta.yaml

@@ -63,12 +63,16 @@ outputs:
         {% endif %}
         - cuda-version ={{ cuda_version }}
         - librmm ={{ minor_version }}
+        - spdlog {{ spdlog_version }}
+        - fmt {{ fmt_version }}
       run:
         - {{ pin_compatible('cuda-version', max_pin='x', min_pin='x') }}
         {% if cuda_major == "11" %}
         - cudatoolkit
         {% endif %}
         - librmm ={{ minor_version }}
+        - spdlog {{ spdlog_version }}
+        - fmt {{ fmt_version }}
     about:
       home: https://rapids.ai/
       license: Apache-2.0

conda/recipes/raft-ann-bench-cpu/conda_build_config.yaml

@@ -18,3 +18,9 @@ h5py_version:
 
 nlohmann_json_version:
   - ">=3.11.2"
+
+spdlog_version:
+  - ">=1.11.0,<1.12"
+
+fmt_version:
+  - ">=9.1.0,<10"

conda/recipes/raft-ann-bench-cpu/meta.yaml

@@ -48,6 +48,8 @@ requirements:
     - glog {{ glog_version }}
     - matplotlib
     - nlohmann_json {{ nlohmann_json_version }}
+    - spdlog {{ spdlog_version }}
+    - fmt {{ fmt_version }}
     - python
     - pyyaml
     - pandas

cpp/bench/ann/src/raft/raft_ivf_flat_wrapper.h

@@ -78,7 +78,7 @@ class RaftIvfFlatGpu : public ANN<T> {
   AlgoProperty get_preference() const override
   {
     AlgoProperty property;
-    property.dataset_memory_type = MemoryType::Device;
+    property.dataset_memory_type = MemoryType::HostMmap;
     property.query_memory_type   = MemoryType::Device;
     return property;
   }

cpp/include/raft/neighbors/detail/ivf_flat_build.cuh

@@ -120,7 +120,8 @@ RAFT_KERNEL build_index_kernel(const LabelT* labels,
                                uint32_t* list_sizes_ptr,
                                IdxT n_rows,
                                uint32_t dim,
-                               uint32_t veclen)
+                               uint32_t veclen,
+                               IdxT batch_offset = 0)
 {
   const IdxT i = IdxT(blockDim.x) * IdxT(blockIdx.x) + threadIdx.x;
   if (i >= n_rows) { return; }
@@ -131,7 +132,7 @@ RAFT_KERNEL build_index_kernel(const LabelT* labels,
   auto* list_data  = list_data_ptrs[list_id];
 
   // Record the source vector id in the index
-  list_index[inlist_id] = source_ixs == nullptr ? i : source_ixs[i];
+  list_index[inlist_id] = source_ixs == nullptr ? i + batch_offset : source_ixs[i];
 
   // The data is written in interleaved groups of `index::kGroupSize` vectors
   using interleaved_group = Pow2<kIndexGroupSize>;
@@ -180,16 +181,33 @@ void extend(raft::resources const& handle,
 
   auto new_labels = raft::make_device_vector<LabelT, IdxT>(handle, n_rows);
   raft::cluster::kmeans_balanced_params kmeans_params;
-  kmeans_params.metric  = index->metric();
-  auto new_vectors_view = raft::make_device_matrix_view<const T, IdxT>(new_vectors, n_rows, dim);
+  kmeans_params.metric = index->metric();
   auto orig_centroids_view =
     raft::make_device_matrix_view<const float, IdxT>(index->centers().data_handle(), n_lists, dim);
-  raft::cluster::kmeans_balanced::predict(handle,
-                                          kmeans_params,
-                                          new_vectors_view,
-                                          orig_centroids_view,
-                                          new_labels.view(),
-                                          utils::mapping<float>{});
+  // Calculate the batch size for the input data if it's not accessible directly from the device
+  constexpr size_t kReasonableMaxBatchSize = 65536;
+  size_t max_batch_size                    = std::min<size_t>(n_rows, kReasonableMaxBatchSize);
+
+  // Predict the cluster labels for the new data, in batches if necessary
+  utils::batch_load_iterator<T> vec_batches(new_vectors,
+                                            n_rows,
+                                            index->dim(),
+                                            max_batch_size,
+                                            stream,
+                                            resource::get_workspace_resource(handle));
+
+  for (const auto& batch : vec_batches) {
+    auto batch_data_view =
+      raft::make_device_matrix_view<const T, IdxT>(batch.data(), batch.size(), index->dim());
+    auto batch_labels_view = raft::make_device_vector_view<LabelT, IdxT>(
+      new_labels.data_handle() + batch.offset(), batch.size());
+    raft::cluster::kmeans_balanced::predict(handle,
+                                            kmeans_params,
+                                            batch_data_view,
+                                            orig_centroids_view,
+                                            batch_labels_view,
+                                            utils::mapping<float>{});
+  }
 
   auto* list_sizes_ptr    = index->list_sizes().data_handle();
   auto old_list_sizes_dev = raft::make_device_vector<uint32_t, IdxT>(handle, n_lists);
@@ -202,14 +220,19 @@ void extend(raft::resources const& handle,
     auto list_sizes_view =
       raft::make_device_vector_view<std::remove_pointer_t<decltype(list_sizes_ptr)>, IdxT>(
         list_sizes_ptr, n_lists);
-    auto const_labels_view = make_const_mdspan(new_labels.view());
-    raft::cluster::kmeans_balanced::helpers::calc_centers_and_sizes(handle,
-                                                                    new_vectors_view,
-                                                                    const_labels_view,
-                                                                    centroids_view,
-                                                                    list_sizes_view,
-                                                                    false,
-                                                                    utils::mapping<float>{});
+    for (const auto& batch : vec_batches) {
+      auto batch_data_view =
+        raft::make_device_matrix_view<const T, IdxT>(batch.data(), batch.size(), index->dim());
+      auto batch_labels_view = raft::make_device_vector_view<const LabelT, IdxT>(
+        new_labels.data_handle() + batch.offset(), batch.size());
+      raft::cluster::kmeans_balanced::helpers::calc_centers_and_sizes(handle,
+                                                                      batch_data_view,
+                                                                      batch_labels_view,
+                                                                      centroids_view,
+                                                                      list_sizes_view,
+                                                                      false,
+                                                                      utils::mapping<float>{});
+    }
   } else {
     raft::stats::histogram<uint32_t, IdxT>(raft::stats::HistTypeAuto,
                                            reinterpret_cast<int32_t*>(list_sizes_ptr),
@@ -244,20 +267,39 @@ void extend(raft::resources const& handle,
   // we'll rebuild the `list_sizes_ptr` in the following kernel, using it as an atomic counter.
   raft::copy(list_sizes_ptr, old_list_sizes_dev.data_handle(), n_lists, stream);
 
-  // Kernel to insert the new vectors
-  const dim3 block_dim(256);
-  const dim3 grid_dim(raft::ceildiv<IdxT>(n_rows, block_dim.x));
-  build_index_kernel<<<grid_dim, block_dim, 0, stream>>>(new_labels.data_handle(),
-                                                         new_vectors,
-                                                         new_indices,
-                                                         index->data_ptrs().data_handle(),
-                                                         index->inds_ptrs().data_handle(),
-                                                         list_sizes_ptr,
-                                                         n_rows,
-                                                         dim,
-                                                         index->veclen());
-  RAFT_CUDA_TRY(cudaPeekAtLastError());
-
+  utils::batch_load_iterator<IdxT> vec_indices(
+    new_indices, n_rows, 1, max_batch_size, stream, resource::get_workspace_resource(handle));
+  utils::batch_load_iterator<IdxT> idx_batch = vec_indices.begin();
+  size_t next_report_offset                  = 0;
+  size_t d_report_offset                     = n_rows * 5 / 100;
+  for (const auto& batch : vec_batches) {
+    auto batch_data_view =
+      raft::make_device_matrix_view<const T, IdxT>(batch.data(), batch.size(), index->dim());
+    // Kernel to insert the new vectors
+    const dim3 block_dim(256);
+    const dim3 grid_dim(raft::ceildiv<IdxT>(batch.size(), block_dim.x));
+    build_index_kernel<T, IdxT, LabelT>
+      <<<grid_dim, block_dim, 0, stream>>>(new_labels.data_handle() + batch.offset(),
+                                           batch_data_view.data_handle(),
+                                           idx_batch->data(),
+                                           index->data_ptrs().data_handle(),
+                                           index->inds_ptrs().data_handle(),
+                                           list_sizes_ptr,
+                                           batch.size(),
+                                           dim,
+                                           index->veclen(),
+                                           batch.offset());
+    RAFT_CUDA_TRY(cudaPeekAtLastError());
+
+    if (batch.offset() > next_report_offset) {
+      float progress = batch.offset() * 100.0f / n_rows;
+      RAFT_LOG_DEBUG("ivf_flat::extend added vectors %zu, %6.1f%% complete",
+                     static_cast<size_t>(batch.offset()),
+                     progress);
+      next_report_offset += d_report_offset;
+    }
+    ++idx_batch;
+  }
   // Precompute the centers vector norms for L2Expanded distance
   if (!index->center_norms().has_value()) {
     index->allocate_center_norms(handle);

cpp/include/raft/neighbors/ivf_flat-ext.cuh

@@ -48,6 +48,18 @@ void build(raft::resources const& handle,
            raft::device_matrix_view<const T, IdxT, row_major> dataset,
            raft::neighbors::ivf_flat::index<T, IdxT>& idx) RAFT_EXPLICIT;
 
+template <typename T, typename IdxT>
+auto build(raft::resources const& handle,
+           const index_params& params,
+           raft::host_matrix_view<const T, IdxT, row_major> dataset)
+  -> index<T, IdxT> RAFT_EXPLICIT;
+
+template <typename T, typename IdxT>
+void build(raft::resources const& handle,
+           const index_params& params,
+           raft::host_matrix_view<const T, IdxT, row_major> dataset,
+           raft::neighbors::ivf_flat::index<T, IdxT>& idx) RAFT_EXPLICIT;
+
 template <typename T, typename IdxT>
 auto extend(raft::resources const& handle,
             const index<T, IdxT>& orig_index,
@@ -74,6 +86,19 @@ void extend(raft::resources const& handle,
             std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices,
             index<T, IdxT>* index) RAFT_EXPLICIT;
 
+template <typename T, typename IdxT>
+auto extend(raft::resources const& handle,
+            raft::host_matrix_view<const T, IdxT, row_major> new_vectors,
+            std::optional<raft::host_vector_view<const IdxT, IdxT>> new_indices,
+            const raft::neighbors::ivf_flat::index<T, IdxT>& orig_index)
+  -> raft::neighbors::ivf_flat::index<T, IdxT> RAFT_EXPLICIT;
+
+template <typename T, typename IdxT>
+void extend(raft::resources const& handle,
+            raft::host_matrix_view<const T, IdxT, row_major> new_vectors,
+            std::optional<raft::host_vector_view<const IdxT, IdxT>> new_indices,
+            index<T, IdxT>* index) RAFT_EXPLICIT;
+
 template <typename T, typename IdxT, typename IvfSampleFilterT>
 void search_with_filtering(raft::resources const& handle,
                            const search_params& params,
@@ -137,6 +162,18 @@ void search(raft::resources const& handle,
     raft::resources const& handle,                                \
     const raft::neighbors::ivf_flat::index_params& params,        \
     raft::device_matrix_view<const T, IdxT, row_major> dataset,   \
+    raft::neighbors::ivf_flat::index<T, IdxT>& idx);              \
+                                                                  \
+  extern template auto raft::neighbors::ivf_flat::build<T, IdxT>( \
+    raft::resources const& handle,                                \
+    const raft::neighbors::ivf_flat::index_params& params,        \
+    raft::host_matrix_view<const T, IdxT, row_major> dataset)     \
+    ->raft::neighbors::ivf_flat::index<T, IdxT>;                  \
+                                                                  \
+  extern template void raft::neighbors::ivf_flat::build<T, IdxT>( \
+    raft::resources const& handle,                                \
+    const raft::neighbors::ivf_flat::index_params& params,        \
+    raft::host_matrix_view<const T, IdxT, row_major> dataset,     \
     raft::neighbors::ivf_flat::index<T, IdxT>& idx);
 
 instantiate_raft_neighbors_ivf_flat_build(float, int64_t);
@@ -171,7 +208,20 @@ instantiate_raft_neighbors_ivf_flat_build(uint8_t, int64_t);
     raft::resources const& handle,                                         \
     raft::device_matrix_view<const T, IdxT, row_major> new_vectors,        \
     std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices, \
-    raft::neighbors::ivf_flat::index<T, IdxT>* index);
+    raft::neighbors::ivf_flat::index<T, IdxT>* index);                     \
+                                                                           \
+  extern template void raft::neighbors::ivf_flat::extend<T, IdxT>(         \
+    raft::resources const& handle,                                         \
+    raft::host_matrix_view<const T, IdxT, row_major> new_vectors,          \
+    std::optional<raft::host_vector_view<const IdxT, IdxT>> new_indices,   \
+    raft::neighbors::ivf_flat::index<T, IdxT>* index);                     \
+                                                                           \
+  extern template auto raft::neighbors::ivf_flat::extend<T, IdxT>(         \
+    const raft::resources& handle,                                         \
+    raft::host_matrix_view<const T, IdxT, row_major> new_vectors,          \
+    std::optional<raft::host_vector_view<const IdxT, IdxT>> new_indices,   \
+    const raft::neighbors::ivf_flat::index<T, IdxT>& idx)                  \
+    ->raft::neighbors::ivf_flat::index<T, IdxT>;
 
 instantiate_raft_neighbors_ivf_flat_extend(float, int64_t);
 instantiate_raft_neighbors_ivf_flat_extend(int8_t, int64_t);