Optimize transfer/compute overlap in out-of-core KMeans

c/include/cuvs/core/c_api.h

@@ -211,6 +211,16 @@ CUVS_EXPORT cuvsError_t cuvsMultiGpuResourcesDestroy(cuvsResources_t res);
  * @return cuvsError_t
  */
 CUVS_EXPORT cuvsError_t cuvsMultiGpuResourcesSetMemoryPool(cuvsResources_t res, int percent_of_free_memory);
+
+/**
+ * @brief Set a CUDA stream pool on all devices managed by the multi-GPU resources
+ *
+ * @param[in] res cuvsResources_t opaque C handle for multi-GPU resources
+ * @param[in] num_streams Number of CUDA streams in each device's pool
+ * @return cuvsError_t
+ */
+CUVS_EXPORT cuvsError_t cuvsMultiGpuResourcesSetStreamPool(cuvsResources_t res,
+                                                           size_t num_streams);
 /** @} */
 
 /**

c/src/core/c_api.cpp

@@ -9,12 +9,16 @@
 #include <raft/core/device_resources_snmg.hpp>
 #include <raft/core/memory_tracking_resources.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resource/cuda_stream_pool.hpp>
 #include <raft/core/resource/device_id.hpp>
 #include <raft/core/resource/device_memory_resource.hpp>
+#include <raft/core/resource/multi_gpu.hpp>
 #include <raft/core/resource/resource_types.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/util/cudart_utils.hpp>
 #include <rapids_logger/logger.hpp>
+#include <rmm/cuda_device.hpp>
+#include <rmm/cuda_stream_pool.hpp>
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/mr/cuda_async_memory_resource.hpp>
 #include <rmm/mr/cuda_memory_resource.hpp>
@@ -132,6 +136,24 @@ extern "C" cuvsError_t cuvsMultiGpuResourcesSetMemoryPool(cuvsResources_t res,
   });
 }
 
+extern "C" cuvsError_t cuvsMultiGpuResourcesSetStreamPool(cuvsResources_t res,
+                                                          size_t num_streams)
+{
+  return cuvs::core::translate_exceptions([=] {
+    RAFT_EXPECTS(num_streams > 0, "num_streams must be greater than zero");
+    auto res_ptr = reinterpret_cast<raft::device_resources_snmg*>(res);
+    RAFT_EXPECTS(res_ptr != nullptr, "res must not be NULL");
+
+    auto& device_resources = raft::resource::get_multi_gpu_resource(*res_ptr);
+    for (auto& device_resource : device_resources) {
+      rmm::cuda_set_device_raii device_guard{
+        rmm::cuda_device_id{raft::resource::get_device_id(device_resource)}};
+      raft::resource::set_cuda_stream_pool(
+        device_resource, std::make_shared<rmm::cuda_stream_pool>(num_streams));
+    }
+  });
+}
+
 extern "C" cuvsError_t cuvsStreamSet(cuvsResources_t res, cudaStream_t stream)
 {
   return cuvs::core::translate_exceptions([=] {

cpp/src/cluster/detail/kmeans.cuh

@@ -5,7 +5,7 @@
 #pragma once
 
 #include "../../core/nvtx.hpp"
-#include "../../neighbors/detail/ann_utils.cuh"
+#include "kmeans_batch_loader.cuh"
 #include "kmeans_common.cuh"
 
 #include <cuvs/cluster/kmeans.hpp>
@@ -24,6 +24,8 @@
 #include <raft/core/pinned_mdarray.hpp>
 #include <raft/core/pinned_mdspan.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resource/cuda_stream_pool.hpp>
+#include <raft/core/resource/device_memory_resource.hpp>
 #include <raft/core/resource/thrust_policy.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/linalg/map.cuh>
@@ -686,25 +688,49 @@ void kmeans_fit(
 
   auto minClusterAndDistance = raft::make_device_vector<raft::KeyValuePair<IndexT, DataT>, IndexT>(
     handle, device_buffer_samples);
-  auto L2NormBatch       = raft::make_device_vector<DataT, IndexT>(handle, device_buffer_samples);
+  const IndexT l2_norm_size = data_on_device ? n_samples : device_buffer_samples;
+  auto L2NormBatch          = raft::make_device_vector<DataT, IndexT>(handle, l2_norm_size);
   auto batch_weights_buf = raft::make_device_vector<DataT, IndexT>(handle, device_buffer_samples);
   rmm::device_uvector<DataT> L2NormBuf_OR_DistBuf(0, stream);
 
   auto centroid_sums      = raft::make_device_matrix<DataT, IndexT>(handle, n_clusters, n_features);
   auto weight_per_cluster = raft::make_device_vector<DataT, IndexT>(handle, n_clusters);
   auto clustering_cost    = raft::make_device_scalar<DataT>(handle, DataT{0});
+  auto batch_inertia      = raft::make_device_scalar<DataT>(handle, DataT{0});
+  auto batch_cost         = raft::make_device_scalar<DataT>(handle, DataT{0});
+  auto h_inertia          = raft::make_pinned_scalar<DataT>(handle, DataT{0});
 
   rmm::device_uvector<char> batch_workspace(device_buffer_samples, stream);
 
-  auto data_batches = cuvs::spatial::knn::detail::utils::make_batch_load_iterator<DataT>(
-    handle, X.data_handle(), n_samples, n_features, device_buffer_samples, stream);
+  auto batch_mr          = data_on_device ? raft::resource::get_workspace_resource_ref(handle)
+                                          : raft::resource::get_large_workspace_resource_ref(handle);
+  auto batch_copy_stream = raft::resource::get_cuda_stream(handle);
+  if constexpr (!data_on_device) {
+    if (handle.has_resource_factory(raft::resource::resource_type::CUDA_STREAM_POOL) &&
+        raft::resource::get_stream_pool_size(handle) >= 1) {
+      batch_copy_stream = raft::resource::get_stream_from_stream_pool(handle);
+    }
+  }
+
+  kmeans_batch_loader<DataT, IndexT, data_on_device> data_batches(handle,
+                                                                  X.data_handle(),
+                                                                  n_samples,
+                                                                  n_features,
+                                                                  device_buffer_samples,
+                                                                  batch_copy_stream,
+                                                                  batch_mr);
   // Host-path weight batches: only materialized when weights are provided and
   // the data resides on host
-  std::optional<cuvs::spatial::knn::detail::utils::batch_load_iterator_dyn<DataT>> weight_batches;
+  std::optional<kmeans_batch_loader<DataT, IndexT, false>> weight_batches;
   if constexpr (!data_on_device) {
     if (weight_ptr != nullptr) {
-      weight_batches = cuvs::spatial::knn::detail::utils::make_batch_load_iterator<DataT>(
-        handle, weight_ptr, n_samples, IndexT{1}, device_buffer_samples, stream);
+      weight_batches.emplace(handle,
+                             weight_ptr,
+                             n_samples,
+                             IndexT{1},
+                             device_buffer_samples,
+                             batch_copy_stream,
+                             batch_mr);
     } else {
       raft::matrix::fill(handle, batch_weights_buf.view(), DataT{1});
     }
@@ -758,6 +784,18 @@ void kmeans_fit(
     }
   };
 
+  auto prefetch_batch = [&](std::size_t batch_pos) {
+    (void)data_batches.prefetch(batch_pos);
+    if (weight_batches.has_value()) { (void)weight_batches->prefetch(batch_pos); }
+  };
+
+  bool input_pipeline_started = false;
+  auto start_input_pipeline   = [&] {
+    if (input_pipeline_started) { return; }
+    if (data_batches.num_batches() > 0) { prefetch_batch(0); }
+    input_pipeline_started = true;
+  };
+
   RAFT_LOG_DEBUG(
     "KMeans.fit: n_samples=%zu, n_features=%zu, n_clusters=%d, device_buffer_samples=%zu",
     static_cast<size_t>(n_samples),
@@ -767,10 +805,6 @@ void kmeans_fit(
 
   bool need_compute_norms = metric == cuvs::distance::DistanceType::L2Expanded ||
                             metric == cuvs::distance::DistanceType::L2SqrtExpanded;
-  auto h_norm_cache = raft::make_pinned_vector<DataT, IndexT>(
-    handle, (need_compute_norms && !data_on_device) ? n_samples : 0);
-  bool norms_cached = false;
-
   auto compute_batch_norms = [&](const DataT* batch_ptr, IndexT batch_size) {
     auto batch_view =
       raft::make_device_matrix_view<const DataT, IndexT>(batch_ptr, batch_size, n_features);
@@ -830,53 +864,44 @@ void kmeans_fit(
       raft::matrix::fill(handle, weight_per_cluster.view(), DataT{0});
       raft::matrix::fill(handle, clustering_cost.view(), DataT{0});
 
+      // Complete iteration setup before starting the cold pipeline, so no potentially blocking
+      // CUDA setup remains between the first transfer and its first consumer.
+      start_input_pipeline();
+
       auto centroids_const = raft::make_device_matrix_view<const DataT, IndexT>(
         cur_centroids_ptr, n_clusters, n_features);
       auto new_centroids_view =
         raft::make_device_matrix_view<DataT, IndexT>(new_centroids_ptr, n_clusters, n_features);
 
-      data_batches.reset();
-      using wt_iter_t = cuvs::spatial::knn::detail::utils::batch_load_iterator_dyn<DataT>;
-      std::optional<wt_iter_t> wt_it;
-      if (weight_batches.has_value()) {
-        weight_batches->reset();
-        wt_it = weight_batches->begin();
-      }
-      for (const auto& data_batch : data_batches) {
-        IndexT cur_batch_size = static_cast<IndexT>(data_batch.size());
-        const DataT* wt_data  = nullptr;
-        if (wt_it.has_value()) {
-          wt_data = (**wt_it).data();
-          ++(*wt_it);
+      for (std::size_t batch_pos = 0; batch_pos < data_batches.num_batches(); ++batch_pos) {
+        const auto data_batch = data_batches.acquire(batch_pos);
+        std::optional<kmeans_batch<DataT>> weight_batch;
+        if (weight_batches.has_value()) {
+          weight_batch.emplace(weight_batches->acquire(batch_pos));
         }
 
+        IndexT cur_batch_size = static_cast<IndexT>(data_batch.size());
+        const DataT* wt_data  = weight_batch.has_value() ? weight_batch->data() : nullptr;
+
         auto batch_data_view = raft::make_device_matrix_view<const DataT, IndexT>(
           data_batch.data(), cur_batch_size, n_features);
         auto batch_weights_view =
           cur_batch_weights(static_cast<IndexT>(data_batch.offset()), wt_data, cur_batch_size);
-
         auto minCAD_view = raft::make_device_vector_view<raft::KeyValuePair<IndexT, DataT>, IndexT>(
           minClusterAndDistance.data_handle(), cur_batch_size);
 
         if constexpr (!data_on_device) {
-          if (need_compute_norms) {
-            if (!norms_cached) {
-              compute_batch_norms(data_batch.data(), cur_batch_size);
-              raft::copy(h_norm_cache.data_handle() + data_batch.offset(),
-                         L2NormBatch.data_handle(),
-                         cur_batch_size,
-                         stream);
-            } else {
-              raft::copy(L2NormBatch.data_handle(),
-                         h_norm_cache.data_handle() + data_batch.offset(),
-                         cur_batch_size,
-                         stream);
-            }
-          }
+          if (need_compute_norms) { compute_batch_norms(data_batch.data(), cur_batch_size); }
         }
 
+        // An already-full pipeline makes this a no-op. During cold fill, submit the first real
+        // consumer before making the second H2D eligible, so CUDA can dispatch both at batch-ready.
+        prefetch_batch((batch_pos + 1) % data_batches.num_batches());
+
+        const auto l2_norm_offset =
+          data_on_device ? static_cast<IndexT>(data_batch.offset()) : IndexT{0};
         auto l2_const_view = raft::make_device_vector_view<const DataT, IndexT>(
-          L2NormBatch.data_handle(), cur_batch_size);
+          L2NormBatch.data_handle() + l2_norm_offset, cur_batch_size);
 
         process_batch<DataT, IndexT>(handle,
                                      batch_data_view,
@@ -892,9 +917,15 @@ void kmeans_fit(
                                      centroid_sums.view(),
                                      weight_per_cluster.view(),
                                      clustering_cost.view(),
-                                     batch_workspace);
+                                     batch_workspace,
+                                     batch_cost.view());
+
+        // The slot is reusable only after every batch consumer above has been submitted. Refill it
+        // with the batch two positions ahead; modulo arithmetic naturally crosses pass boundaries.
+        const auto next_batch_pos = (batch_pos + 2) % data_batches.num_batches();
+        data_batches.recycle(data_batch, next_batch_pos);
+        if (weight_batch.has_value()) { weight_batches->recycle(*weight_batch, next_batch_pos); }
       }
-      if (need_compute_norms) { norms_cached = true; }
 
       finalize_centroids<DataT, IndexT>(handle,
                                         raft::make_const_mdspan(centroid_sums.view()),
@@ -927,46 +958,89 @@ void kmeans_fit(
       raft::copy(handle,
                  raft::make_pinned_scalar_view(h_done_flag.data_handle()),
                  raft::make_device_scalar_view<const int>(d_done_flag.data_handle()));
+      // The next pass's first two input batches are already in flight. The compute stream still
+      // serializes centroid finalization and convergence before it can consume them.
     }
 
     {
       auto centroids_const = raft::make_device_matrix_view<const DataT, IndexT>(
         cur_centroids_ptr, n_clusters, n_features);
 
-      iter_inertia = DataT{0};
-      data_batches.reset();
-      using wt_iter_t = cuvs::spatial::knn::detail::utils::batch_load_iterator_dyn<DataT>;
-      std::optional<wt_iter_t> wt_it;
-      if (weight_batches.has_value()) {
-        weight_batches->reset();
-        wt_it = weight_batches->begin();
-      }
-      for (const auto& data_batch : data_batches) {
-        IndexT cur_batch_size = static_cast<IndexT>(data_batch.size());
-        const DataT* wt_data  = nullptr;
-        if (wt_it.has_value()) {
-          wt_data = (**wt_it).data();
-          ++(*wt_it);
+      raft::matrix::fill(handle, clustering_cost.view(), DataT{0});
+      start_input_pipeline();
+      for (std::size_t batch_pos = 0; batch_pos < data_batches.num_batches(); ++batch_pos) {
+        const auto data_batch = data_batches.acquire(batch_pos);
+        std::optional<kmeans_batch<DataT>> weight_batch;
+        if (weight_batches.has_value()) {
+          weight_batch.emplace(weight_batches->acquire(batch_pos));
         }
 
+        IndexT cur_batch_size = static_cast<IndexT>(data_batch.size());
+        const DataT* wt_data  = weight_batch.has_value() ? weight_batch->data() : nullptr;
+
         auto batch_data_view = raft::make_device_matrix_view<const DataT, IndexT>(
           data_batch.data(), cur_batch_size, n_features);
-
         std::optional<raft::device_vector_view<const DataT, IndexT>> batch_sw = std::nullopt;
         if (weight_ptr != nullptr) {
           batch_sw =
             cur_batch_weights(static_cast<IndexT>(data_batch.offset()), wt_data, cur_batch_size);
         }
+        compute_batch_norms(data_batch.data(), cur_batch_size);
+        if (batch_pos + 1 < data_batches.num_batches() || seed_iter + 1 < n_init) {
+          prefetch_batch((batch_pos + 1) % data_batches.num_batches());
+        }
 
-        DataT batch_cost = DataT{0};
-        cuvs::cluster::kmeans::cluster_cost(handle,
-                                            batch_data_view,
-                                            centroids_const,
-                                            raft::make_host_scalar_view(&batch_cost),
-                                            batch_sw);
-
-        iter_inertia += batch_cost;
+        auto l2_norm_view = raft::make_device_vector_view<const DataT, IndexT>(
+          L2NormBatch.data_handle(), cur_batch_size);
+        auto min_cad_view =
+          raft::make_device_vector_view<raft::KeyValuePair<IndexT, DataT>, IndexT>(
+            minClusterAndDistance.data_handle(), cur_batch_size);
+
+        minClusterAndDistanceCompute<DataT, IndexT>(handle,
+                                                    batch_data_view,
+                                                    centroids_const,
+                                                    min_cad_view,
+                                                    l2_norm_view,
+                                                    L2NormBuf_OR_DistBuf,
+                                                    cuvs::distance::DistanceType::L2Expanded,
+                                                    iter_params.batch_samples,
+                                                    iter_params.batch_centroids,
+                                                    ws);
+        if (batch_sw.has_value()) {
+          raft::linalg::map(
+            handle,
+            min_cad_view,
+            [] __device__(raft::KeyValuePair<IndexT, DataT> pair, DataT weight) {
+              pair.value *= weight;
+              return pair;
+            },
+            raft::make_const_mdspan(min_cad_view),
+            batch_sw.value());
+        }
+        computeClusterCost(
+          handle, min_cad_view, ws, batch_inertia.view(), raft::value_op{}, raft::add_op{});
+        raft::linalg::add(clustering_cost.data_handle(),
+                          clustering_cost.data_handle(),
+                          batch_inertia.data_handle(),
+                          1,
+                          stream);
+
+        const bool needs_future_batch =
+          batch_pos + 2 < data_batches.num_batches() || seed_iter + 1 < n_init;
+        if (needs_future_batch) {
+          const auto next_batch_pos = (batch_pos + 2) % data_batches.num_batches();
+          data_batches.recycle(data_batch, next_batch_pos);
+          if (weight_batch.has_value()) { weight_batches->recycle(*weight_batch, next_batch_pos); }
+        } else {
+          data_batches.release(data_batch);
+          if (weight_batch.has_value()) { weight_batches->release(*weight_batch); }
+        }
       }
+      raft::copy(handle,
+                 raft::make_pinned_scalar_view(h_inertia.data_handle()),
+                 raft::make_device_scalar_view<const DataT>(clustering_cost.data_handle()));
+      raft::resource::sync_stream(handle);
+      iter_inertia = *h_inertia.data_handle();
     }
 
     if (iter_inertia < inertia[0]) {