Enable OOC KMeans transfer/compute overlap for pageable input

c/include/cuvs/cluster/kmeans.h

@@ -105,11 +105,18 @@ struct cuvsKMeansParams {
 
   /**
    * Number of samples to process per GPU batch for the batched (host-data) API.
-   * When set to 0, defaults to n_samples (process all at once).
+   * Default: 0, which uses n_samples (process all at once). For a single host
+   * input matrix, this creates one full-dataset input buffer on the GPU, even
+   * when a stream pool is configured. Set a positive value less than n_samples
+   * to stream that matrix in smaller batches.
    * Multiple batches use one input buffer without a stream pool and two when
-   * an auxiliary stream enables transfer/compute overlap. Each input buffer
-   * uses about device_buffer_samples * n_features * sizeof(value_type) bytes,
-   * excluding algorithm workspaces and optional sample-weight buffers.
+   * an auxiliary stream enables transfer/compute overlap. No stream pool is
+   * configured by default, so transfers and compute are serialized. Call
+   * cuvsResourcesSetStreamPool(res, 1) to enable overlap for multiple batches.
+   * Each input buffer uses about
+   * device_buffer_samples * n_features * sizeof(value_type) bytes (substitute
+   * n_samples when device_buffer_samples is 0), excluding algorithm workspaces
+   * and optional sample-weight buffers.
    */
   int64_t device_buffer_samples;
 

cpp/include/cuvs/cluster/kmeans.hpp

@@ -148,16 +148,22 @@ struct params : base_params {
    * Only used by the batched (host-data) code path and ignored by
    * device-data overloads.
    *
+   * Default: 0. For a single host input matrix, this creates one full-dataset input
+   * buffer on the GPU, even when an auxiliary stream is configured. To stream that
+   * matrix in smaller batches, set a positive value less than n_samples.
+   *
    * Inputs spanning multiple batches are double-buffered when the handle has an auxiliary
-   * stream, and use one buffer otherwise. Budget about
+   * stream, and use one buffer otherwise. No auxiliary stream pool is configured by
+   * default, so transfers and compute are serialized. Configure a pool with at least
+   * one auxiliary stream to enable transfer/compute overlap. Budget about
    * `device_buffer_samples * n_features * sizeof(value_type)` bytes per input buffer,
    * plus algorithm workspaces. Sample weights require the same number of additional buffers
-   * with `device_buffer_samples * sizeof(value_type)` bytes each.
+   * with `device_buffer_samples * sizeof(value_type)` bytes each. In these estimates,
+   * substitute n_samples for device_buffer_samples when it is 0.
    *
    * In multi-GPU mode this is a per-rank batch size: each rank processes up
    * to this many local samples per batch, clamped to that rank's local sample
    * count. This is ignored by device-data overloads.
-   * Default: 0 (process all data at once).
    */
   int64_t device_buffer_samples = 0;
 };
@@ -261,8 +267,10 @@ enum class kmeans_type { KMeans = 0, KMeansBalanced = 1 };
  * Multiple batches are double-buffered when an auxiliary stream is available.
  * Without a stream pool, the fit uses one input buffer and remains correct but
  * serialized. One auxiliary stream is sufficient to overlap transfer and compute.
- * Pinned host memory is crucial for OOC performance: ordinary pageable or
- * unregistered `mmap`-backed input degrades throughput rapidly.
+ * Pageable and unregistered `mmap`-backed input is read through
+ * one batch-sized pinned host staging buffer per device buffer. The complete input is
+ * never registered or copied into pinned memory. Already pinned input is transferred directly.
+ * Overlap depends on the host reader supplying batches fast enough for the GPU.
  *
  * A pool memory resource is optional but recommended, especially for repeated
  * fits. Configure both pools before `fit` and before first using the handle's
@@ -1747,9 +1755,11 @@ void cluster_cost(
  * `device_buffer_samples` is ignored and each local partition is processed in full.
  * Host partitions spanning multiple batches are double-buffered when an auxiliary stream
  * is available. Each rank needs one auxiliary stream for transfer/compute overlap;
- * otherwise execution uses one input buffer and is correct but serialized. Pinned host memory is
- * crucial: pageable or unregistered `mmap`-backed input degrades throughput rapidly. A per-device
- * memory pool is also recommended.
+ * otherwise execution uses one input buffer and is correct but serialized. Pageable and
+ * unregistered `mmap`-backed partitions use bounded pinned host staging buffers;
+ * already pinned partitions are transferred directly. Overlap depends on the
+ * host reader supplying batches fast enough for the GPU. A per-device memory pool is also
+ * recommended.
  * With `raft::device_resources_snmg`, configure these before `fit` with
  * `handle.set_stream_pool(1)` and `handle.set_memory_pool(percent)`; both calls
  * apply the setting to every device managed by the handle. With independently

cpp/src/cluster/detail/kmeans_batch_loader.cuh

@@ -8,6 +8,7 @@
 #include <raft/core/device_mdspan.hpp>
 #include <raft/core/error.hpp>
 #include <raft/core/host_mdspan.hpp>
+#include <raft/core/pinned_mdarray.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/util/cuda_utils.cuh>
@@ -23,6 +24,7 @@
 
 #include <algorithm>
 #include <cstddef>
+#include <cstring>
 #include <optional>
 #include <utility>
 #include <vector>
@@ -88,7 +90,8 @@ class kmeans_batch {
  *
  * The device specialization is a zero-copy view. The host specialization below owns the cyclic
  * H2D pipeline used by out-of-core KMeans: two buffers when prefetching is enabled, or one buffer
- * for serialized execution.
+ * for serialized execution. Pageable host partitions use batch-sized pinned staging buffers;
+ * the input itself is never registered or copied in full.
  */
 template <typename DataT, typename IndexT, bool DataOnDevice>
 class kmeans_batch_loader;
@@ -189,6 +192,8 @@ class kmeans_batch_loader<DataT, IndexT, false> {
   {
     for (std::size_t partition = 0; partition < partitions.size(); ++partition) {
       append_batches(partitions[partition], partition);
+      pageable_partitions_.push_back(partitions[partition].extent(0) > 0 &&
+                                     is_pageable(partitions[partition].data_handle()));
     }
     if (batches_.empty()) { return; }
 
@@ -204,6 +209,16 @@ class kmeans_batch_loader<DataT, IndexT, false> {
       buffer_1_.resize(max_batch_elements, copy_stream_);
       buffer_ptrs_[1] = buffer_1_.data();
     }
+
+    if (std::any_of(
+          pageable_partitions_.begin(), pageable_partitions_.end(), [](bool v) { return v; })) {
+      // Only a batch-sized host allocation per device slot. Already pinned/registered inputs
+      // retain the original direct H2D path without staging memory.
+      for (int slot = 0; slot < num_slots(); ++slot) {
+        host_buffers_[slot].emplace(
+          raft::make_pinned_vector<DataT, std::size_t>(res, max_batch_elements));
+      }
+    }
   }
 
   kmeans_batch_loader(kmeans_batch_loader const&)                    = delete;
@@ -213,8 +228,11 @@ class kmeans_batch_loader<DataT, IndexT, false> {
 
   ~kmeans_batch_loader() noexcept
   {
-    if (!batches_.empty()) { raft::resource::sync_stream(*res_); }
-    raft::resource::sync_stream(*res_, copy_stream_);
+    // Drain both streams before events and staging buffers are destroyed, including on exceptions.
+    if (!batches_.empty()) {
+      raft::resource::sync_stream(*res_);
+      raft::resource::sync_stream(*res_, copy_stream_);
+    }
     for (auto event : events_) {
       if (event != nullptr) { RAFT_CUDA_TRY_NO_THROW(cudaEventDestroy(event)); }
     }
@@ -322,7 +340,7 @@ class kmeans_batch_loader<DataT, IndexT, false> {
     if (states_[slot] == slot_state::reusable) {
       RAFT_CUDA_TRY(cudaStreamWaitEvent(copy_stream_.get(), reusable_[slot], 0));
     }
-    queue_h2d(buffer_ptrs_[slot], pos);
+    queue_h2d(slot, pos);
     positions_[slot] = pos;
     if (ready_[slot] == nullptr) { ready_[slot] = make_event(); }
     // cudaStreamWaitEvent captures the latest record at the time the wait is submitted, so this
@@ -348,13 +366,41 @@ class kmeans_batch_loader<DataT, IndexT, false> {
     return slot;
   }
 
-  void queue_h2d(DataT* dst, std::size_t pos)
+  void queue_h2d(int slot, std::size_t pos)
   {
     auto const& batch = batches_[pos];
-    raft::copy(dst,
-               batch.input.data_handle() + batch.descriptor.offset() * batch.input.extent(1),
-               batch.descriptor.size() * batch.input.extent(1),
-               copy_stream_);
+    auto const* source =
+      batch.input.data_handle() + batch.descriptor.offset() * batch.input.extent(1);
+    const auto elements = batch.descriptor.size() * batch.input.extent(1);
+    const bool pageable = pageable_partitions_[batch.descriptor.partition()];
+    if (pageable) {
+      // The host slot can be refilled after its previous H2D, while its device slot is computing.
+      if (host_reusable_[slot] != nullptr) {
+        RAFT_CUDA_TRY(cudaEventSynchronize(host_reusable_[slot]));
+      } else {
+        host_reusable_[slot] = make_event();
+      }
+      std::memcpy(host_buffers_[slot]->data_handle(), source, elements * sizeof(DataT));
+      source = host_buffers_[slot]->data_handle();
+    }
+    // Pinned H2D submission returns without waiting for the device-slot reuse event. Direct
+    // pageable cudaMemcpyAsync may block here and prevent the next batch's compute being queued.
+    raft::copy(buffer_ptrs_[slot], source, elements, copy_stream_);
+    if (pageable) { RAFT_CUDA_TRY(cudaEventRecord(host_reusable_[slot], copy_stream_.get())); }
+  }
+
+  static bool is_pageable(DataT const* ptr)
+  {
+    cudaPointerAttributes attributes{};
+    auto status = cudaPointerGetAttributes(&attributes, ptr);
+    if (status == cudaErrorInvalidValue) {
+      // Older CUDA runtimes report unregistered host pointers as invalid rather than returning
+      // cudaMemoryTypeUnregistered. Clear only this expected error before subsequent CUDA calls.
+      static_cast<void>(cudaGetLastError());
+      return true;
+    }
+    RAFT_CUDA_TRY(status);
+    return attributes.type == cudaMemoryTypeUnregistered;
   }
 
   raft::resources const* res_ = nullptr;
@@ -371,6 +417,10 @@ class kmeans_batch_loader<DataT, IndexT, false> {
   cudaEvent_t ready_[2]    = {nullptr, nullptr};
   cudaEvent_t reusable_[2] = {nullptr, nullptr};
   std::vector<cudaEvent_t> events_;
+  std::vector<bool> pageable_partitions_;
+  std::optional<raft::pinned_vector<DataT, std::size_t>> host_buffers_[2];
+  // Separate from ready_: mixed pinned partitions can reuse ready_ without using host_buffers_.
+  cudaEvent_t host_reusable_[2] = {nullptr, nullptr};
 };
 
 }  // namespace cuvs::cluster::kmeans::detail

cpp/tests/CMakeLists.txt

@@ -138,6 +138,7 @@ ConfigureTest(
        cluster/kmeans_find_k.cu
        cluster/gmm.cu
        cluster/kmeans_predict_batching.cu
+       cluster/kmeans_batch_loader.cu
        cluster/linkage.cu
        cluster/connect_knn.cu
        cluster/spectral.cu