From 95e372975fce6313abe0134014418a8d7a6a6f32 Mon Sep 17 00:00:00 2001 From: viclafargue Date: Tue, 1 Sep 2026 10:38:40 +0000 Subject: [PATCH 1/3] Optimize out-of-core KMeans overlap --- cpp/src/cluster/detail/kmeans.cuh | 176 +++++++----- .../cluster/detail/kmeans_batch_loader.cuh | 258 ++++++++++++++++++ cpp/src/cluster/detail/kmeans_common.cuh | 52 +++- cpp/src/cluster/detail/kmeans_mg.cuh | 4 +- cpp/tests/cluster/kmeans.cu | 64 +++++ 5 files changed, 485 insertions(+), 69 deletions(-) create mode 100644 cpp/src/cluster/detail/kmeans_batch_loader.cuh diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index e3ffb4a439..81f1cf2884 100644 --- a/cpp/src/cluster/detail/kmeans.cuh +++ b/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 @@ -24,6 +24,8 @@ #include #include #include +#include +#include #include #include #include @@ -39,6 +41,7 @@ #include #include +#include #include #include @@ -686,25 +689,57 @@ void kmeans_fit( auto minClusterAndDistance = raft::make_device_vector, IndexT>( handle, device_buffer_samples); - auto L2NormBatch = raft::make_device_vector(handle, device_buffer_samples); + auto minClusterDistance = raft::make_device_vector(handle, device_buffer_samples); + const IndexT l2_norm_size = data_on_device ? n_samples : device_buffer_samples; + auto L2NormBatch = raft::make_device_vector(handle, l2_norm_size); auto batch_weights_buf = raft::make_device_vector(handle, device_buffer_samples); rmm::device_uvector L2NormBuf_OR_DistBuf(0, stream); auto centroid_sums = raft::make_device_matrix(handle, n_clusters, n_features); auto weight_per_cluster = raft::make_device_vector(handle, n_clusters); auto clustering_cost = raft::make_device_scalar(handle, DataT{0}); + auto batch_inertia = raft::make_device_scalar(handle, DataT{0}); + auto batch_cost = raft::make_device_scalar(handle, DataT{0}); + auto h_inertia = raft::make_pinned_scalar(handle, DataT{0}); rmm::device_uvector batch_workspace(device_buffer_samples, stream); - auto data_batches = cuvs::spatial::knn::detail::utils::make_batch_load_iterator( - 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); + std::optional owned_batch_copy_stream; + if constexpr (!data_on_device) { + // Host KMeans must pipeline even when a minimal caller (including the Python binding) did not + // install a RAFT stream pool. Keep the stream owner before the iterators so their destructors + // drain all event dependencies before the stream itself is destroyed. + 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); + } else { + owned_batch_copy_stream.emplace(rmm::cuda_stream::flags::non_blocking); + batch_copy_stream = owned_batch_copy_stream->view(); + } + } + + kmeans_batch_loader 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> weight_batches; + std::optional> weight_batches; if constexpr (!data_on_device) { if (weight_ptr != nullptr) { - weight_batches = cuvs::spatial::knn::detail::utils::make_batch_load_iterator( - 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}); } @@ -767,10 +802,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( - 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(batch_ptr, batch_size, n_features); @@ -835,48 +866,41 @@ void kmeans_fit( auto new_centroids_view = raft::make_device_matrix_view(new_centroids_ptr, n_clusters, n_features); - data_batches.reset(); - using wt_iter_t = cuvs::spatial::knn::detail::utils::batch_load_iterator_dyn; - std::optional wt_it; - if (weight_batches.has_value()) { - weight_batches->reset(); - wt_it = weight_batches->begin(); - } - for (const auto& data_batch : data_batches) { + data_batches.prime(); + if (weight_batches.has_value()) { weight_batches->prime(); } + for (std::size_t batch_pos = 0; batch_pos < data_batches.num_batches(); ++batch_pos) { + const auto data_batch = data_batches.load(batch_pos); IndexT cur_batch_size = static_cast(data_batch.size()); - const DataT* wt_data = nullptr; - if (wt_it.has_value()) { - wt_data = (**wt_it).data(); - ++(*wt_it); - } + const DataT* wt_data = + weight_batches.has_value() ? weight_batches->load(batch_pos).data() : nullptr; auto batch_data_view = raft::make_device_matrix_view( data_batch.data(), cur_batch_size, n_features); auto batch_weights_view = cur_batch_weights(static_cast(data_batch.offset()), wt_data, cur_batch_size); - auto minCAD_view = raft::make_device_vector_view, 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); } + } + // Submit the bandwidth-bound row norm before the next bulk H2D, then submit the remaining + // assignment work while that copy is in flight. This keeps the copy call timely without + // letting it get ahead of (and starve) the norm kernel. + const bool is_last_batch = + data_batch.offset() + data_batch.size() == static_cast(n_samples); + if (is_last_batch) { + data_batches.prime(); + if (weight_batches.has_value()) { weight_batches->prime(); } + } else { + data_batches.prefetch(batch_pos + 1); + if (weight_batches.has_value()) { weight_batches->prefetch(batch_pos + 1); } } + const auto l2_norm_offset = + data_on_device ? static_cast(data_batch.offset()) : IndexT{0}; auto l2_const_view = raft::make_device_vector_view( - L2NormBatch.data_handle(), cur_batch_size); + L2NormBatch.data_handle() + l2_norm_offset, cur_batch_size); process_batch(handle, batch_data_view, @@ -892,9 +916,9 @@ void kmeans_fit( centroid_sums.view(), weight_per_cluster.view(), clustering_cost.view(), - batch_workspace); + batch_workspace, + batch_cost.view()); } - if (need_compute_norms) { norms_cached = true; } finalize_centroids(handle, raft::make_const_mdspan(centroid_sums.view()), @@ -927,46 +951,68 @@ void kmeans_fit( raft::copy(handle, raft::make_pinned_scalar_view(h_done_flag.data_handle()), raft::make_device_scalar_view(d_done_flag.data_handle())); + // Batch 0 was queued before last-batch compute. Retire that slot only after convergence is + // submitted, then queue batch 1 so neither transfer delays these latency-critical kernels. + data_batches.prime_second_batch(); + if (weight_batches.has_value()) { weight_batches->prime_second_batch(); } } { auto centroids_const = raft::make_device_matrix_view( 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; - std::optional wt_it; - if (weight_batches.has_value()) { - weight_batches->reset(); - wt_it = weight_batches->begin(); - } - for (const auto& data_batch : data_batches) { + raft::matrix::fill(handle, clustering_cost.view(), DataT{0}); + data_batches.prime(); + if (weight_batches.has_value()) { weight_batches->prime(); } + for (std::size_t batch_pos = 0; batch_pos < data_batches.num_batches(); ++batch_pos) { + const auto data_batch = data_batches.load(batch_pos); IndexT cur_batch_size = static_cast(data_batch.size()); - const DataT* wt_data = nullptr; - if (wt_it.has_value()) { - wt_data = (**wt_it).data(); - ++(*wt_it); - } + const DataT* wt_data = + weight_batches.has_value() ? weight_batches->load(batch_pos).data() : nullptr; auto batch_data_view = raft::make_device_matrix_view( data_batch.data(), cur_batch_size, n_features); - std::optional> batch_sw = std::nullopt; if (weight_ptr != nullptr) { batch_sw = cur_batch_weights(static_cast(data_batch.offset()), wt_data, cur_batch_size); } + compute_batch_norms(data_batch.data(), cur_batch_size); + const bool is_last_batch = + data_batch.offset() + data_batch.size() == static_cast(n_samples); + if (!is_last_batch) { + data_batches.prefetch(batch_pos + 1); + if (weight_batches.has_value()) { weight_batches->prefetch(batch_pos + 1); } + } - 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(L2NormBatch.data_handle(), cur_batch_size); + auto min_distance_view = raft::make_device_vector_view( + minClusterDistance.data_handle(), cur_batch_size); + + cluster_cost(handle, + batch_data_view, + centroids_const, + min_distance_view, + l2_norm_view, + L2NormBuf_OR_DistBuf, + cuvs::distance::DistanceType::L2Expanded, + iter_params.batch_samples, + iter_params.batch_centroids, + ws, + batch_inertia.view(), + batch_sw); + raft::linalg::add(clustering_cost.data_handle(), + clustering_cost.data_handle(), + batch_inertia.data_handle(), + 1, + stream); } + raft::copy(handle, + raft::make_pinned_scalar_view(h_inertia.data_handle()), + raft::make_device_scalar_view(clustering_cost.data_handle())); + raft::resource::sync_stream(handle); + iter_inertia = *h_inertia.data_handle(); } if (iter_inertia < inertia[0]) { diff --git a/cpp/src/cluster/detail/kmeans_batch_loader.cuh b/cpp/src/cluster/detail/kmeans_batch_loader.cuh new file mode 100644 index 0000000000..f1ee5d6253 --- /dev/null +++ b/cpp/src/cluster/detail/kmeans_batch_loader.cuh @@ -0,0 +1,258 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +#pragma once + +#include +#include +#include +#include +#include + +#include +#include +#include + +#include + +#include +#include +#include +#include +#include + +namespace cuvs::cluster::kmeans::detail { + +/** A contiguous KMeans input batch accessible from the main CUDA stream. */ +template +class kmeans_batch { + public: + [[nodiscard]] auto data() const noexcept -> DataT const* { return data_; } + [[nodiscard]] auto size() const noexcept -> std::size_t { return size_; } + [[nodiscard]] auto offset() const noexcept -> std::size_t { return offset_; } + + private: + template + friend class kmeans_batch_loader; + + kmeans_batch(DataT const* data, std::size_t size, std::size_t offset) + : data_(data), size_(size), offset_(offset) + { + } + + DataT const* data_ = nullptr; + std::size_t size_ = 0; + std::size_t offset_ = 0; +}; + +/** + * Read-only batch loader used only by KMeans. + * + * The device specialization is a zero-copy view. The host specialization below owns the + * two-buffer, cyclic H2D pipeline needed by out-of-core KMeans. + */ +template +class kmeans_batch_loader; + +template +class kmeans_batch_loader { + public: + kmeans_batch_loader(raft::resources const&, + DataT const* source, + IndexT n_rows, + IndexT row_width, + IndexT batch_size, + rmm::cuda_stream_view, + rmm::device_async_resource_ref) + : source_(source), + n_rows_(static_cast(n_rows)), + row_width_(static_cast(row_width)), + batch_size_(std::min(static_cast(batch_size), + std::max(n_rows_, 1))), + n_batches_(n_rows_ == 0 ? 0 : raft::div_rounding_up_safe(n_rows_, batch_size_)) + { + } + + [[nodiscard]] auto num_batches() const noexcept -> std::size_t { return n_batches_; } + void prime() noexcept {} + void prefetch(std::size_t) noexcept {} + void prime_second_batch() noexcept {} + + [[nodiscard]] auto load(std::size_t pos) const -> kmeans_batch + { + RAFT_EXPECTS(pos < n_batches_, "KMeans batch position is out of range"); + const auto offset = pos * batch_size_; + const auto size = std::min(batch_size_, n_rows_ - offset); + return {source_ + offset * row_width_, size, offset}; + } + + private: + DataT const* source_ = nullptr; + std::size_t n_rows_ = 0; + std::size_t row_width_ = 0; + std::size_t batch_size_ = 0; + std::size_t n_batches_ = 0; +}; + +template +class kmeans_batch_loader { + public: + kmeans_batch_loader(raft::resources const& res, + DataT const* source, + IndexT n_rows, + IndexT row_width, + IndexT batch_size, + rmm::cuda_stream_view copy_stream, + rmm::device_async_resource_ref mr) + : res_(&res), + source_(source), + n_rows_(static_cast(n_rows)), + row_width_(static_cast(row_width)), + batch_size_(std::min(static_cast(batch_size), + std::max(n_rows_, 1))), + n_batches_(n_rows_ == 0 ? 0 : raft::div_rounding_up_safe(n_rows_, batch_size_)), + copy_stream_(copy_stream), + buffer_0_(0, copy_stream, mr), + buffer_1_(0, copy_stream, mr) + { + if (n_rows_ == 0 || source_ == nullptr) { return; } + + buffer_0_.resize(row_width_ * batch_size_, copy_stream_); + current_ptr_ = buffer_0_.data(); + if (n_batches_ > 1) { + buffer_1_.resize(row_width_ * batch_size_, copy_stream_); + prefetch_ptr_ = buffer_1_.data(); + } + } + + kmeans_batch_loader(kmeans_batch_loader const&) = delete; + auto operator=(kmeans_batch_loader const&) -> kmeans_batch_loader& = delete; + kmeans_batch_loader(kmeans_batch_loader&&) = delete; + auto operator=(kmeans_batch_loader&&) -> kmeans_batch_loader& = delete; + + ~kmeans_batch_loader() noexcept + { + if (source_ != nullptr) { + RAFT_CUDA_TRY_NO_THROW(cudaStreamSynchronize(raft::resource::get_cuda_stream(*res_))); + } + RAFT_CUDA_TRY_NO_THROW(cudaStreamSynchronize(copy_stream_)); + for (auto event : events_) { + if (event != nullptr) { RAFT_CUDA_TRY_NO_THROW(cudaEventDestroy(event)); } + } + } + + [[nodiscard]] auto num_batches() const noexcept -> std::size_t { return n_batches_; } + + /** Stage batch zero unless it is already active or staged by the previous pass. */ + void prime() + { + if (n_batches_ <= 1) { return; } + const bool batch_zero_active = current_pos_.has_value() && *current_pos_ == 0; + const bool batch_zero_staged = prefetch_pos_.has_value() && *prefetch_pos_ == 0; + if (!batch_zero_active && !batch_zero_staged) { prefetch(0); } + } + + /** Stage a future batch into the slot not currently consumed by KMeans. */ + void prefetch(std::size_t pos) + { + if (n_batches_ <= 1 || pos >= n_batches_ || source_ == nullptr) { return; } + if (prefetch_pos_.has_value() && *prefetch_pos_ == pos) { return; } + + const int prefetch_slot = 1 - current_slot_; + if (kernel_done_[prefetch_slot] != nullptr) { + RAFT_CUDA_TRY(cudaStreamWaitEvent(copy_stream_, kernel_done_[prefetch_slot], 0)); + } + + queue_h2d(prefetch_ptr_, pos); + prefetch_pos_ = pos; + h2d_done_[prefetch_slot] = make_event(); + RAFT_CUDA_TRY(cudaEventRecord(h2d_done_[prefetch_slot], copy_stream_)); + } + + /** + * Activate cyclic batch zero after pass-boundary kernels have been submitted, then stage batch + * one into the slot retired by the previous pass's last batch. + */ + void prime_second_batch() + { + prime(); + if (n_batches_ < 2) { return; } + (void)load(0); + prefetch(1); + } + + /** Make a staged batch visible to kernels on the main stream. */ + [[nodiscard]] auto load(std::size_t pos) -> kmeans_batch + { + RAFT_EXPECTS(pos < n_batches_, "KMeans batch position is out of range"); + if (!current_pos_.has_value() || *current_pos_ != pos) { + if (prefetch_pos_.has_value() && *prefetch_pos_ == pos) { + const int retired_slot = current_slot_; + std::swap(current_ptr_, prefetch_ptr_); + current_slot_ = 1 - current_slot_; + prefetch_pos_.reset(); + + kernel_done_[retired_slot] = make_event(); + RAFT_CUDA_TRY( + cudaEventRecord(kernel_done_[retired_slot], raft::resource::get_cuda_stream(*res_))); + RAFT_CUDA_TRY( + cudaStreamWaitEvent(raft::resource::get_cuda_stream(*res_), h2d_done_[current_slot_], 0)); + } else { + // A one-batch input has nothing to overlap. Stage it once, then reuse it for every pass. + RAFT_EXPECTS(n_batches_ == 1, "KMeans attempted to load a batch that was not prefetched"); + queue_h2d(current_ptr_, pos); + copy_stream_.synchronize(); + } + current_pos_ = pos; + } + + const auto offset = pos * batch_size_; + const auto size = std::min(batch_size_, n_rows_ - offset); + return {current_ptr_, size, offset}; + } + + private: + [[nodiscard]] auto make_event() -> cudaEvent_t + { + cudaEvent_t event = nullptr; + RAFT_CUDA_TRY(cudaEventCreateWithFlags(&event, cudaEventDisableTiming)); + try { + events_.push_back(event); + } catch (...) { + RAFT_CUDA_TRY_NO_THROW(cudaEventDestroy(event)); + throw; + } + return event; + } + + void queue_h2d(DataT* dst, std::size_t pos) + { + const auto offset = pos * batch_size_; + const auto rows = std::min(batch_size_, n_rows_ - offset); + const auto bytes = rows * row_width_ * sizeof(DataT); + RAFT_CUDA_TRY(cudaMemcpyAsync( + dst, source_ + offset * row_width_, bytes, cudaMemcpyHostToDevice, copy_stream_)); + } + + raft::resources const* res_ = nullptr; + DataT const* source_ = nullptr; + std::size_t n_rows_ = 0; + std::size_t row_width_ = 0; + std::size_t batch_size_ = 0; + std::size_t n_batches_ = 0; + rmm::cuda_stream_view copy_stream_; + rmm::device_uvector buffer_0_; + rmm::device_uvector buffer_1_; + DataT* current_ptr_ = nullptr; + DataT* prefetch_ptr_ = nullptr; + int current_slot_ = 0; + std::optional current_pos_; + std::optional prefetch_pos_; + cudaEvent_t h2d_done_[2] = {nullptr, nullptr}; + cudaEvent_t kernel_done_[2] = {nullptr, nullptr}; + std::vector events_; +}; + +} // namespace cuvs::cluster::kmeans::detail diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index ab3ef0a05a..bd4c487eec 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -446,6 +446,50 @@ EXTERN_TEMPLATE_MIN_CLUSTER_DISTANCE(double, int) #undef EXTERN_TEMPLATE_MIN_CLUSTER_DISTANCE +/** + * @brief Compute the optionally weighted sum of distances to the nearest centroid. + * + * Unlike minClusterAndDistanceCompute, this path does not calculate cluster labels. + */ +template +void cluster_cost( + raft::resources const& handle, + raft::device_matrix_view X, + raft::device_matrix_view centroids, + raft::device_vector_view min_cluster_distance, + raft::device_vector_view l2_norm_x, + rmm::device_uvector& l2_norm_or_distance_buffer, + cuvs::distance::DistanceType metric, + int batch_samples, + int batch_centroids, + rmm::device_uvector& workspace, + raft::device_scalar_view cost, + std::optional> sample_weight = std::nullopt) +{ + auto centroids_mutable = raft::make_device_matrix_view( + const_cast(centroids.data_handle()), centroids.extent(0), centroids.extent(1)); + minClusterDistanceCompute(handle, + X, + centroids_mutable, + min_cluster_distance, + l2_norm_x, + l2_norm_or_distance_buffer, + metric, + batch_samples, + batch_centroids, + workspace); + + if (sample_weight.has_value()) { + raft::linalg::map(handle, + min_cluster_distance, + raft::mul_op{}, + raft::make_const_mdspan(min_cluster_distance), + sample_weight.value()); + } + computeClusterCost( + handle, min_cluster_distance, workspace, cost, raft::identity_op{}, raft::add_op{}); +} + template void countSamplesInCluster(raft::resources const& handle, const cuvs::cluster::kmeans::params& params, @@ -683,6 +727,8 @@ __device__ void check_convergence(raft::device_scalar_view clusteri * @param[inout] centroid_sums Running weighted sums [n_clusters x n_features] (added into) * @param[inout] weight_per_cluster Running weight counts [n_clusters] (added into) * @param[inout] clustering_cost Running cost scalar (device) (added into) + * @param[out] batch_cost Scratch scalar (device) for this batch's cost. Owned by the + * caller so that streaming loops make no allocation per batch. */ template void process_batch( @@ -700,7 +746,8 @@ void process_batch( raft::device_matrix_view centroid_sums, raft::device_vector_view weight_per_cluster, raft::device_scalar_view clustering_cost, - rmm::device_uvector& batch_workspace) + rmm::device_uvector& batch_workspace, + raft::device_scalar_view batch_cost) { cudaStream_t stream = raft::resource::get_cuda_stream(handle); @@ -742,9 +789,8 @@ void process_batch( raft::make_const_mdspan(minClusterAndDistance), batch_weights); - auto batch_cost = raft::make_device_scalar(handle, DataT{0}); computeClusterCost( - handle, minClusterAndDistance, workspace, batch_cost.view(), raft::value_op{}, raft::add_op{}); + handle, minClusterAndDistance, workspace, batch_cost, raft::value_op{}, raft::add_op{}); raft::linalg::add(clustering_cost.data_handle(), clustering_cost.data_handle(), batch_cost.data_handle(), diff --git a/cpp/src/cluster/detail/kmeans_mg.cuh b/cpp/src/cluster/detail/kmeans_mg.cuh index dbe2c23039..c6871020ca 100644 --- a/cpp/src/cluster/detail/kmeans_mg.cuh +++ b/cpp/src/cluster/detail/kmeans_mg.cuh @@ -212,6 +212,7 @@ void mnmg_fit( auto clustering_cost = raft::make_device_vector(dev_res, 1); auto batch_clustering_cost = raft::make_device_vector(dev_res, 1); auto sqrd_norm_error_dev = raft::make_device_scalar(dev_res, DataT{0}); + auto batch_cost = raft::make_device_scalar(dev_res, DataT{0}); IndexT alloc_batch_size = device_buffer_samples; auto batch_weights = raft::make_device_vector(dev_res, alloc_batch_size); auto minClusterAndDistance = @@ -467,7 +468,8 @@ void mnmg_fit( centroid_sums.view(), weight_per_cluster.view(), raft::make_device_scalar_view(clustering_cost.data_handle()), - batch_workspace); + batch_workspace, + batch_cost.view()); } } norms_cached = true; diff --git a/cpp/tests/cluster/kmeans.cu b/cpp/tests/cluster/kmeans.cu index 804922a93d..b54ef8e16e 100644 --- a/cpp/tests/cluster/kmeans.cu +++ b/cpp/tests/cluster/kmeans.cu @@ -3,6 +3,7 @@ * SPDX-License-Identifier: Apache-2.0 */ +#include "../../src/cluster/detail/kmeans_batch_loader.cuh" #include "../test_utils.cuh" #include "kmeans_test_blobs.cuh" @@ -11,11 +12,13 @@ #include #include #include +#include #include #include #include #include +#include #include #include @@ -706,4 +709,65 @@ INSTANTIATE_TEST_CASE_P(KmeansFitBatchedTests, KmeansFitBatchedTestD, ::testing::ValuesIn(batched_inputsd2)); +TEST(KmeansBatchLoaderTest, CyclicFourPasses) +{ + constexpr int64_t n_rows = 257; + constexpr int64_t n_cols = 17; + constexpr int64_t batch_size = 64; + constexpr int n_passes = 4; + + raft::resources handle; + rmm::cuda_stream copy_stream(rmm::cuda_stream::flags::non_blocking); + std::vector host_data(n_rows * n_cols); + for (int64_t row = 0; row < n_rows; ++row) { + for (int64_t col = 0; col < n_cols; ++col) { + host_data[row * n_cols + col] = row * n_cols + col; + } + } + + cluster::kmeans::detail::kmeans_batch_loader loader( + handle, + host_data.data(), + n_rows, + n_cols, + batch_size, + copy_stream, + raft::resource::get_workspace_resource_ref(handle)); + auto device_readback = + raft::make_device_vector(handle, n_passes * n_rows * n_cols); + + loader.prime(); + for (int pass = 0; pass < n_passes; ++pass) { + for (std::size_t pos = 0; pos < loader.num_batches(); ++pos) { + const auto batch = loader.load(pos); + const bool last = pos + 1 == loader.num_batches(); + if (last && pass + 1 < n_passes) { + loader.prime(); + } else if (!last) { + loader.prefetch(pos + 1); + } + + const auto output_offset = + (static_cast(pass) * n_rows + batch.offset()) * n_cols; + raft::copy(device_readback.data_handle() + output_offset, + batch.data(), + batch.size() * n_cols, + raft::resource::get_cuda_stream(handle)); + if (last && pass + 1 < n_passes) { loader.prime_second_batch(); } + } + } + + std::vector readback(device_readback.size()); + raft::copy(readback.data(), + device_readback.data_handle(), + device_readback.size(), + raft::resource::get_cuda_stream(handle)); + raft::resource::sync_stream(handle); + for (int pass = 0; pass < n_passes; ++pass) { + for (std::size_t i = 0; i < host_data.size(); ++i) { + EXPECT_EQ(readback[static_cast(pass) * host_data.size() + i], host_data[i]); + } + } +} + } // namespace cuvs From 9c16817eae06423f81d46216a8c3db88a7918cb6 Mon Sep 17 00:00:00 2001 From: viclafargue Date: Tue, 1 Sep 2026 13:50:02 +0000 Subject: [PATCH 2/3] Refine out-of-core KMeans batch pipelining --- cpp/src/cluster/detail/kmeans.cuh | 134 +++++++++------ .../cluster/detail/kmeans_batch_loader.cuh | 161 ++++++++++-------- cpp/tests/cluster/kmeans.cu | 22 +-- 3 files changed, 189 insertions(+), 128 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index 81f1cf2884..5fa92e9ade 100644 --- a/cpp/src/cluster/detail/kmeans.cuh +++ b/cpp/src/cluster/detail/kmeans.cuh @@ -689,7 +689,6 @@ void kmeans_fit( auto minClusterAndDistance = raft::make_device_vector, IndexT>( handle, device_buffer_samples); - auto minClusterDistance = raft::make_device_vector(handle, device_buffer_samples); const IndexT l2_norm_size = data_on_device ? n_samples : device_buffer_samples; auto L2NormBatch = raft::make_device_vector(handle, l2_norm_size); auto batch_weights_buf = raft::make_device_vector(handle, device_buffer_samples); @@ -793,6 +792,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(n_samples), @@ -861,18 +872,24 @@ 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( cur_centroids_ptr, n_clusters, n_features); auto new_centroids_view = raft::make_device_matrix_view(new_centroids_ptr, n_clusters, n_features); - data_batches.prime(); - if (weight_batches.has_value()) { weight_batches->prime(); } for (std::size_t batch_pos = 0; batch_pos < data_batches.num_batches(); ++batch_pos) { - const auto data_batch = data_batches.load(batch_pos); + const auto data_batch = data_batches.acquire(batch_pos); + std::optional> weight_batch; + if (weight_batches.has_value()) { + weight_batch.emplace(weight_batches->acquire(batch_pos)); + } + IndexT cur_batch_size = static_cast(data_batch.size()); - const DataT* wt_data = - weight_batches.has_value() ? weight_batches->load(batch_pos).data() : nullptr; + const DataT* wt_data = weight_batch.has_value() ? weight_batch->data() : nullptr; auto batch_data_view = raft::make_device_matrix_view( data_batch.data(), cur_batch_size, n_features); @@ -884,18 +901,10 @@ void kmeans_fit( if constexpr (!data_on_device) { if (need_compute_norms) { compute_batch_norms(data_batch.data(), cur_batch_size); } } - // Submit the bandwidth-bound row norm before the next bulk H2D, then submit the remaining - // assignment work while that copy is in flight. This keeps the copy call timely without - // letting it get ahead of (and starve) the norm kernel. - const bool is_last_batch = - data_batch.offset() + data_batch.size() == static_cast(n_samples); - if (is_last_batch) { - data_batches.prime(); - if (weight_batches.has_value()) { weight_batches->prime(); } - } else { - data_batches.prefetch(batch_pos + 1); - if (weight_batches.has_value()) { weight_batches->prefetch(batch_pos + 1); } - } + + // 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(data_batch.offset()) : IndexT{0}; @@ -918,6 +927,12 @@ void kmeans_fit( clustering_cost.view(), 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); } } finalize_centroids(handle, @@ -951,10 +966,8 @@ void kmeans_fit( raft::copy(handle, raft::make_pinned_scalar_view(h_done_flag.data_handle()), raft::make_device_scalar_view(d_done_flag.data_handle())); - // Batch 0 was queued before last-batch compute. Retire that slot only after convergence is - // submitted, then queue batch 1 so neither transfer delays these latency-critical kernels. - data_batches.prime_second_batch(); - if (weight_batches.has_value()) { weight_batches->prime_second_batch(); } + // 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. } { @@ -962,13 +975,16 @@ void kmeans_fit( cur_centroids_ptr, n_clusters, n_features); raft::matrix::fill(handle, clustering_cost.view(), DataT{0}); - data_batches.prime(); - if (weight_batches.has_value()) { weight_batches->prime(); } + start_input_pipeline(); for (std::size_t batch_pos = 0; batch_pos < data_batches.num_batches(); ++batch_pos) { - const auto data_batch = data_batches.load(batch_pos); + const auto data_batch = data_batches.acquire(batch_pos); + std::optional> weight_batch; + if (weight_batches.has_value()) { + weight_batch.emplace(weight_batches->acquire(batch_pos)); + } + IndexT cur_batch_size = static_cast(data_batch.size()); - const DataT* wt_data = - weight_batches.has_value() ? weight_batches->load(batch_pos).data() : nullptr; + const DataT* wt_data = weight_batch.has_value() ? weight_batch->data() : nullptr; auto batch_data_view = raft::make_device_matrix_view( data_batch.data(), cur_batch_size, n_features); @@ -978,35 +994,55 @@ void kmeans_fit( cur_batch_weights(static_cast(data_batch.offset()), wt_data, cur_batch_size); } compute_batch_norms(data_batch.data(), cur_batch_size); - const bool is_last_batch = - data_batch.offset() + data_batch.size() == static_cast(n_samples); - if (!is_last_batch) { - data_batches.prefetch(batch_pos + 1); - if (weight_batches.has_value()) { weight_batches->prefetch(batch_pos + 1); } + if (batch_pos + 1 < data_batches.num_batches() || seed_iter + 1 < n_init) { + prefetch_batch((batch_pos + 1) % data_batches.num_batches()); } - auto l2_norm_view = - raft::make_device_vector_view(L2NormBatch.data_handle(), cur_batch_size); - auto min_distance_view = raft::make_device_vector_view( - minClusterDistance.data_handle(), cur_batch_size); - - cluster_cost(handle, - batch_data_view, - centroids_const, - min_distance_view, - l2_norm_view, - L2NormBuf_OR_DistBuf, - cuvs::distance::DistanceType::L2Expanded, - iter_params.batch_samples, - iter_params.batch_centroids, - ws, - batch_inertia.view(), - batch_sw); + auto l2_norm_view = raft::make_device_vector_view( + L2NormBatch.data_handle(), cur_batch_size); + auto min_cad_view = + raft::make_device_vector_view, IndexT>( + minClusterAndDistance.data_handle(), cur_batch_size); + + minClusterAndDistanceCompute(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 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()), diff --git a/cpp/src/cluster/detail/kmeans_batch_loader.cuh b/cpp/src/cluster/detail/kmeans_batch_loader.cuh index f1ee5d6253..eb71b1781d 100644 --- a/cpp/src/cluster/detail/kmeans_batch_loader.cuh +++ b/cpp/src/cluster/detail/kmeans_batch_loader.cuh @@ -37,14 +37,15 @@ class kmeans_batch { template friend class kmeans_batch_loader; - kmeans_batch(DataT const* data, std::size_t size, std::size_t offset) - : data_(data), size_(size), offset_(offset) + kmeans_batch(DataT const* data, std::size_t size, std::size_t offset, int slot) + : data_(data), size_(size), offset_(offset), slot_(slot) { } DataT const* data_ = nullptr; std::size_t size_ = 0; std::size_t offset_ = 0; + int slot_ = 0; }; /** @@ -76,16 +77,16 @@ class kmeans_batch_loader { } [[nodiscard]] auto num_batches() const noexcept -> std::size_t { return n_batches_; } - void prime() noexcept {} void prefetch(std::size_t) noexcept {} - void prime_second_batch() noexcept {} + void recycle(kmeans_batch const&, std::size_t) noexcept {} + void release(kmeans_batch const&) noexcept {} - [[nodiscard]] auto load(std::size_t pos) const -> kmeans_batch + [[nodiscard]] auto acquire(std::size_t pos) const -> kmeans_batch { RAFT_EXPECTS(pos < n_batches_, "KMeans batch position is out of range"); const auto offset = pos * batch_size_; const auto size = std::min(batch_size_, n_rows_ - offset); - return {source_ + offset * row_width_, size, offset}; + return {source_ + offset * row_width_, size, offset, 0}; } private: @@ -120,10 +121,10 @@ class kmeans_batch_loader { if (n_rows_ == 0 || source_ == nullptr) { return; } buffer_0_.resize(row_width_ * batch_size_, copy_stream_); - current_ptr_ = buffer_0_.data(); + buffer_ptrs_[0] = buffer_0_.data(); if (n_batches_ > 1) { buffer_1_.resize(row_width_ * batch_size_, copy_stream_); - prefetch_ptr_ = buffer_1_.data(); + buffer_ptrs_[1] = buffer_1_.data(); } } @@ -145,75 +146,65 @@ class kmeans_batch_loader { [[nodiscard]] auto num_batches() const noexcept -> std::size_t { return n_batches_; } - /** Stage batch zero unless it is already active or staged by the previous pass. */ - void prime() - { - if (n_batches_ <= 1) { return; } - const bool batch_zero_active = current_pos_.has_value() && *current_pos_ == 0; - const bool batch_zero_staged = prefetch_pos_.has_value() && *prefetch_pos_ == 0; - if (!batch_zero_active && !batch_zero_staged) { prefetch(0); } - } - - /** Stage a future batch into the slot not currently consumed by KMeans. */ + /** Stage a batch into an available slot; do nothing when both slots are occupied. */ void prefetch(std::size_t pos) { - if (n_batches_ <= 1 || pos >= n_batches_ || source_ == nullptr) { return; } - if (prefetch_pos_.has_value() && *prefetch_pos_ == pos) { return; } + RAFT_EXPECTS(pos < n_batches_, "KMeans batch position is out of range"); + if (source_ == nullptr) { return; } - const int prefetch_slot = 1 - current_slot_; - if (kernel_done_[prefetch_slot] != nullptr) { - RAFT_CUDA_TRY(cudaStreamWaitEvent(copy_stream_, kernel_done_[prefetch_slot], 0)); + for (int slot = 0; slot < num_slots(); ++slot) { + if (states_[slot] == slot_state::empty || states_[slot] == slot_state::reusable) { + stage(slot, pos); + return; + } } - - queue_h2d(prefetch_ptr_, pos); - prefetch_pos_ = pos; - h2d_done_[prefetch_slot] = make_event(); - RAFT_CUDA_TRY(cudaEventRecord(h2d_done_[prefetch_slot], copy_stream_)); } - /** - * Activate cyclic batch zero after pass-boundary kernels have been submitted, then stage batch - * one into the slot retired by the previous pass's last batch. - */ - void prime_second_batch() + /** Make a prefetched batch visible to kernels on the main stream. */ + [[nodiscard]] auto acquire(std::size_t pos) -> kmeans_batch { - prime(); - if (n_batches_ < 2) { return; } - (void)load(0); - prefetch(1); + RAFT_EXPECTS(pos < n_batches_, "KMeans batch position is out of range"); + for (int slot = 0; slot < num_slots(); ++slot) { + if (states_[slot] == slot_state::staged && positions_[slot] == pos) { + RAFT_CUDA_TRY(cudaStreamWaitEvent(raft::resource::get_cuda_stream(*res_), ready_[slot], 0)); + states_[slot] = slot_state::acquired; + + const auto offset = pos * batch_size_; + const auto size = std::min(batch_size_, n_rows_ - offset); + return {buffer_ptrs_[slot], size, offset, slot}; + } + } + RAFT_FAIL("KMeans attempted to acquire a batch that was not prefetched"); } - /** Make a staged batch visible to kernels on the main stream. */ - [[nodiscard]] auto load(std::size_t pos) -> kmeans_batch + /** Record completion of a batch, then refill the same slot with a future batch. */ + void recycle(kmeans_batch const& batch, std::size_t next_pos) { - RAFT_EXPECTS(pos < n_batches_, "KMeans batch position is out of range"); - if (!current_pos_.has_value() || *current_pos_ != pos) { - if (prefetch_pos_.has_value() && *prefetch_pos_ == pos) { - const int retired_slot = current_slot_; - std::swap(current_ptr_, prefetch_ptr_); - current_slot_ = 1 - current_slot_; - prefetch_pos_.reset(); - - kernel_done_[retired_slot] = make_event(); - RAFT_CUDA_TRY( - cudaEventRecord(kernel_done_[retired_slot], raft::resource::get_cuda_stream(*res_))); - RAFT_CUDA_TRY( - cudaStreamWaitEvent(raft::resource::get_cuda_stream(*res_), h2d_done_[current_slot_], 0)); - } else { - // A one-batch input has nothing to overlap. Stage it once, then reuse it for every pass. - RAFT_EXPECTS(n_batches_ == 1, "KMeans attempted to load a batch that was not prefetched"); - queue_h2d(current_ptr_, pos); - copy_stream_.synchronize(); - } - current_pos_ = pos; + RAFT_EXPECTS(next_pos < n_batches_, "KMeans batch position is out of range"); + const int slot = validate_acquired(batch); + + // No transfer is needed when the requested future batch is already resident. + if (positions_[slot] == next_pos) { + states_[slot] = slot_state::staged; + return; } - const auto offset = pos * batch_size_; - const auto size = std::min(batch_size_, n_rows_ - offset); - return {current_ptr_, size, offset}; + mark_reusable(slot); + stage(slot, next_pos); + } + + /** Record completion without scheduling another transfer into the slot. */ + void release(kmeans_batch const& batch) + { + const int slot = validate_acquired(batch); + mark_reusable(slot); } private: + enum class slot_state { empty, staged, acquired, reusable }; + + [[nodiscard]] auto num_slots() const noexcept -> int { return n_batches_ > 1 ? 2 : 1; } + [[nodiscard]] auto make_event() -> cudaEvent_t { cudaEvent_t event = nullptr; @@ -227,6 +218,40 @@ class kmeans_batch_loader { return event; } + void stage(int slot, std::size_t pos) + { + RAFT_EXPECTS(states_[slot] == slot_state::empty || states_[slot] == slot_state::reusable, + "KMeans attempted to overwrite an active batch buffer"); + if (states_[slot] == slot_state::reusable) { + RAFT_CUDA_TRY(cudaStreamWaitEvent(copy_stream_, reusable_[slot], 0)); + } + queue_h2d(buffer_ptrs_[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 + // per-slot event can be reused after acquire() has enqueued that wait. + RAFT_CUDA_TRY(cudaEventRecord(ready_[slot], copy_stream_)); + states_[slot] = slot_state::staged; + } + + void mark_reusable(int slot) + { + if (reusable_[slot] == nullptr) { reusable_[slot] = make_event(); } + // The copy stream consumes this generation's record before the event is recorded again. + RAFT_CUDA_TRY(cudaEventRecord(reusable_[slot], raft::resource::get_cuda_stream(*res_))); + states_[slot] = slot_state::reusable; + } + + [[nodiscard]] auto validate_acquired(kmeans_batch const& batch) const -> int + { + const int slot = batch.slot_; + RAFT_EXPECTS(slot >= 0 && slot < num_slots() && states_[slot] == slot_state::acquired && + positions_[slot] == batch.offset() / batch_size_ && + buffer_ptrs_[slot] == batch.data(), + "KMeans attempted to release a batch that is not active"); + return slot; + } + void queue_h2d(DataT* dst, std::size_t pos) { const auto offset = pos * batch_size_; @@ -245,13 +270,11 @@ class kmeans_batch_loader { rmm::cuda_stream_view copy_stream_; rmm::device_uvector buffer_0_; rmm::device_uvector buffer_1_; - DataT* current_ptr_ = nullptr; - DataT* prefetch_ptr_ = nullptr; - int current_slot_ = 0; - std::optional current_pos_; - std::optional prefetch_pos_; - cudaEvent_t h2d_done_[2] = {nullptr, nullptr}; - cudaEvent_t kernel_done_[2] = {nullptr, nullptr}; + DataT* buffer_ptrs_[2] = {nullptr, nullptr}; + std::optional positions_[2]; + slot_state states_[2] = {slot_state::empty, slot_state::empty}; + cudaEvent_t ready_[2] = {nullptr, nullptr}; + cudaEvent_t reusable_[2] = {nullptr, nullptr}; std::vector events_; }; diff --git a/cpp/tests/cluster/kmeans.cu b/cpp/tests/cluster/kmeans.cu index b54ef8e16e..9fcb15ca90 100644 --- a/cpp/tests/cluster/kmeans.cu +++ b/cpp/tests/cluster/kmeans.cu @@ -736,24 +736,26 @@ TEST(KmeansBatchLoaderTest, CyclicFourPasses) auto device_readback = raft::make_device_vector(handle, n_passes * n_rows * n_cols); - loader.prime(); + if (loader.num_batches() > 0) { loader.prefetch(0); } for (int pass = 0; pass < n_passes; ++pass) { for (std::size_t pos = 0; pos < loader.num_batches(); ++pos) { - const auto batch = loader.load(pos); - const bool last = pos + 1 == loader.num_batches(); - if (last && pass + 1 < n_passes) { - loader.prime(); - } else if (!last) { - loader.prefetch(pos + 1); - } - + const auto batch = loader.acquire(pos); const auto output_offset = (static_cast(pass) * n_rows + batch.offset()) * n_cols; raft::copy(device_readback.data_handle() + output_offset, batch.data(), batch.size() * n_cols, raft::resource::get_cuda_stream(handle)); - if (last && pass + 1 < n_passes) { loader.prime_second_batch(); } + + if (pos + 1 < loader.num_batches() || pass + 1 < n_passes) { + loader.prefetch((pos + 1) % loader.num_batches()); + } + const bool needs_future_batch = pos + 2 < loader.num_batches() || pass + 1 < n_passes; + if (needs_future_batch) { + loader.recycle(batch, (pos + 2) % loader.num_batches()); + } else { + loader.release(batch); + } } } From 4b8a5c541223ea83266344a9b4240671ab49ad54 Mon Sep 17 00:00:00 2001 From: viclafargue Date: Wed, 2 Sep 2026 15:27:56 +0000 Subject: [PATCH 3/3] multi-GPU transfer/compute overlap --- .../cluster/detail/kmeans_batch_loader.cuh | 172 ++++++--- cpp/src/cluster/detail/kmeans_mg.cuh | 342 ++++++++++++------ 2 files changed, 346 insertions(+), 168 deletions(-) diff --git a/cpp/src/cluster/detail/kmeans_batch_loader.cuh b/cpp/src/cluster/detail/kmeans_batch_loader.cuh index eb71b1781d..51ef308c4e 100644 --- a/cpp/src/cluster/detail/kmeans_batch_loader.cuh +++ b/cpp/src/cluster/detail/kmeans_batch_loader.cuh @@ -25,6 +25,21 @@ namespace cuvs::cluster::kmeans::detail { +/** One independently-addressed input partition in a logical KMeans batch sequence. */ +template +struct kmeans_input_partition { + DataT const* data; + IndexT size; +}; + +template +struct kmeans_batch_descriptor { + DataT const* source; + std::size_t size; + std::size_t offset; + std::size_t partition; +}; + /** A contiguous KMeans input batch accessible from the main CUDA stream. */ template class kmeans_batch { @@ -32,20 +47,33 @@ class kmeans_batch { [[nodiscard]] auto data() const noexcept -> DataT const* { return data_; } [[nodiscard]] auto size() const noexcept -> std::size_t { return size_; } [[nodiscard]] auto offset() const noexcept -> std::size_t { return offset_; } + [[nodiscard]] auto partition() const noexcept -> std::size_t { return partition_; } private: template friend class kmeans_batch_loader; - kmeans_batch(DataT const* data, std::size_t size, std::size_t offset, int slot) - : data_(data), size_(size), offset_(offset), slot_(slot) + kmeans_batch(DataT const* data, + std::size_t size, + std::size_t offset, + std::size_t partition, + std::size_t position, + int slot) + : data_(data), + size_(size), + offset_(offset), + partition_(partition), + position_(position), + slot_(slot) { } - DataT const* data_ = nullptr; - std::size_t size_ = 0; - std::size_t offset_ = 0; - int slot_ = 0; + DataT const* data_ = nullptr; + std::size_t size_ = 0; + std::size_t offset_ = 0; + std::size_t partition_ = 0; + std::size_t position_ = 0; + int slot_ = 0; }; /** @@ -60,41 +88,64 @@ class kmeans_batch_loader; template class kmeans_batch_loader { public: - kmeans_batch_loader(raft::resources const&, + kmeans_batch_loader(raft::resources const& res, DataT const* source, IndexT n_rows, IndexT row_width, IndexT batch_size, + rmm::cuda_stream_view copy_stream, + rmm::device_async_resource_ref mr) + : kmeans_batch_loader(res, + std::vector>{{source, n_rows}}, + row_width, + batch_size, + copy_stream, + mr) + { + } + + kmeans_batch_loader(raft::resources const&, + std::vector> const& partitions, + IndexT row_width, + IndexT batch_size, rmm::cuda_stream_view, rmm::device_async_resource_ref) - : source_(source), - n_rows_(static_cast(n_rows)), - row_width_(static_cast(row_width)), - batch_size_(std::min(static_cast(batch_size), - std::max(n_rows_, 1))), - n_batches_(n_rows_ == 0 ? 0 : raft::div_rounding_up_safe(n_rows_, batch_size_)) + : row_width_(static_cast(row_width)), + batch_size_(std::max(static_cast(batch_size), 1)) { + for (std::size_t partition = 0; partition < partitions.size(); ++partition) { + append_batches(partitions[partition], partition); + } } - [[nodiscard]] auto num_batches() const noexcept -> std::size_t { return n_batches_; } + [[nodiscard]] auto num_batches() const noexcept -> std::size_t { return batches_.size(); } void prefetch(std::size_t) noexcept {} void recycle(kmeans_batch const&, std::size_t) noexcept {} void release(kmeans_batch const&) noexcept {} [[nodiscard]] auto acquire(std::size_t pos) const -> kmeans_batch { - RAFT_EXPECTS(pos < n_batches_, "KMeans batch position is out of range"); - const auto offset = pos * batch_size_; - const auto size = std::min(batch_size_, n_rows_ - offset); - return {source_ + offset * row_width_, size, offset, 0}; + RAFT_EXPECTS(pos < batches_.size(), "KMeans batch position is out of range"); + auto const& batch = batches_[pos]; + return { + batch.source + batch.offset * row_width_, batch.size, batch.offset, batch.partition, pos, 0}; } private: - DataT const* source_ = nullptr; - std::size_t n_rows_ = 0; + void append_batches(kmeans_input_partition input, std::size_t partition) + { + const auto n_rows = static_cast(input.size); + if (n_rows == 0) { return; } + RAFT_EXPECTS(input.data != nullptr, "non-empty KMeans input partition cannot be null"); + for (std::size_t offset = 0; offset < n_rows; offset += batch_size_) { + const auto size = std::min(batch_size_, n_rows - offset); + batches_.push_back({input.data, size, offset, partition}); + } + } + std::size_t row_width_ = 0; std::size_t batch_size_ = 0; - std::size_t n_batches_ = 0; + std::vector> batches_; }; template @@ -107,23 +158,41 @@ class kmeans_batch_loader { IndexT batch_size, rmm::cuda_stream_view copy_stream, rmm::device_async_resource_ref mr) + : kmeans_batch_loader(res, + std::vector>{{source, n_rows}}, + row_width, + batch_size, + copy_stream, + mr) + { + } + + kmeans_batch_loader(raft::resources const& res, + std::vector> const& partitions, + IndexT row_width, + IndexT batch_size, + rmm::cuda_stream_view copy_stream, + rmm::device_async_resource_ref mr) : res_(&res), - source_(source), - n_rows_(static_cast(n_rows)), row_width_(static_cast(row_width)), - batch_size_(std::min(static_cast(batch_size), - std::max(n_rows_, 1))), - n_batches_(n_rows_ == 0 ? 0 : raft::div_rounding_up_safe(n_rows_, batch_size_)), + batch_size_(std::max(static_cast(batch_size), 1)), copy_stream_(copy_stream), buffer_0_(0, copy_stream, mr), buffer_1_(0, copy_stream, mr) { - if (n_rows_ == 0 || source_ == nullptr) { return; } + for (std::size_t partition = 0; partition < partitions.size(); ++partition) { + append_batches(partitions[partition], partition); + } + if (batches_.empty()) { return; } - buffer_0_.resize(row_width_ * batch_size_, copy_stream_); + std::size_t max_batch_rows = 0; + for (auto const& batch : batches_) { + max_batch_rows = std::max(max_batch_rows, batch.size); + } + buffer_0_.resize(row_width_ * max_batch_rows, copy_stream_); buffer_ptrs_[0] = buffer_0_.data(); - if (n_batches_ > 1) { - buffer_1_.resize(row_width_ * batch_size_, copy_stream_); + if (batches_.size() > 1) { + buffer_1_.resize(row_width_ * max_batch_rows, copy_stream_); buffer_ptrs_[1] = buffer_1_.data(); } } @@ -135,7 +204,7 @@ class kmeans_batch_loader { ~kmeans_batch_loader() noexcept { - if (source_ != nullptr) { + if (!batches_.empty()) { RAFT_CUDA_TRY_NO_THROW(cudaStreamSynchronize(raft::resource::get_cuda_stream(*res_))); } RAFT_CUDA_TRY_NO_THROW(cudaStreamSynchronize(copy_stream_)); @@ -144,13 +213,12 @@ class kmeans_batch_loader { } } - [[nodiscard]] auto num_batches() const noexcept -> std::size_t { return n_batches_; } + [[nodiscard]] auto num_batches() const noexcept -> std::size_t { return batches_.size(); } /** Stage a batch into an available slot; do nothing when both slots are occupied. */ void prefetch(std::size_t pos) { - RAFT_EXPECTS(pos < n_batches_, "KMeans batch position is out of range"); - if (source_ == nullptr) { return; } + RAFT_EXPECTS(pos < batches_.size(), "KMeans batch position is out of range"); for (int slot = 0; slot < num_slots(); ++slot) { if (states_[slot] == slot_state::empty || states_[slot] == slot_state::reusable) { @@ -163,15 +231,14 @@ class kmeans_batch_loader { /** Make a prefetched batch visible to kernels on the main stream. */ [[nodiscard]] auto acquire(std::size_t pos) -> kmeans_batch { - RAFT_EXPECTS(pos < n_batches_, "KMeans batch position is out of range"); + RAFT_EXPECTS(pos < batches_.size(), "KMeans batch position is out of range"); for (int slot = 0; slot < num_slots(); ++slot) { if (states_[slot] == slot_state::staged && positions_[slot] == pos) { RAFT_CUDA_TRY(cudaStreamWaitEvent(raft::resource::get_cuda_stream(*res_), ready_[slot], 0)); states_[slot] = slot_state::acquired; - const auto offset = pos * batch_size_; - const auto size = std::min(batch_size_, n_rows_ - offset); - return {buffer_ptrs_[slot], size, offset, slot}; + auto const& batch = batches_[pos]; + return {buffer_ptrs_[slot], batch.size, batch.offset, batch.partition, pos, slot}; } } RAFT_FAIL("KMeans attempted to acquire a batch that was not prefetched"); @@ -180,7 +247,7 @@ class kmeans_batch_loader { /** Record completion of a batch, then refill the same slot with a future batch. */ void recycle(kmeans_batch const& batch, std::size_t next_pos) { - RAFT_EXPECTS(next_pos < n_batches_, "KMeans batch position is out of range"); + RAFT_EXPECTS(next_pos < batches_.size(), "KMeans batch position is out of range"); const int slot = validate_acquired(batch); // No transfer is needed when the requested future batch is already resident. @@ -203,7 +270,18 @@ class kmeans_batch_loader { private: enum class slot_state { empty, staged, acquired, reusable }; - [[nodiscard]] auto num_slots() const noexcept -> int { return n_batches_ > 1 ? 2 : 1; } + [[nodiscard]] auto num_slots() const noexcept -> int { return batches_.size() > 1 ? 2 : 1; } + + void append_batches(kmeans_input_partition input, std::size_t partition) + { + const auto n_rows = static_cast(input.size); + if (n_rows == 0) { return; } + RAFT_EXPECTS(input.data != nullptr, "non-empty KMeans input partition cannot be null"); + for (std::size_t offset = 0; offset < n_rows; offset += batch_size_) { + const auto size = std::min(batch_size_, n_rows - offset); + batches_.push_back({input.data, size, offset, partition}); + } + } [[nodiscard]] auto make_event() -> cudaEvent_t { @@ -246,27 +324,23 @@ class kmeans_batch_loader { { const int slot = batch.slot_; RAFT_EXPECTS(slot >= 0 && slot < num_slots() && states_[slot] == slot_state::acquired && - positions_[slot] == batch.offset() / batch_size_ && - buffer_ptrs_[slot] == batch.data(), + positions_[slot] == batch.position_ && buffer_ptrs_[slot] == batch.data(), "KMeans attempted to release a batch that is not active"); return slot; } void queue_h2d(DataT* dst, std::size_t pos) { - const auto offset = pos * batch_size_; - const auto rows = std::min(batch_size_, n_rows_ - offset); - const auto bytes = rows * row_width_ * sizeof(DataT); + auto const& batch = batches_[pos]; + const auto bytes = batch.size * row_width_ * sizeof(DataT); RAFT_CUDA_TRY(cudaMemcpyAsync( - dst, source_ + offset * row_width_, bytes, cudaMemcpyHostToDevice, copy_stream_)); + dst, batch.source + batch.offset * row_width_, bytes, cudaMemcpyHostToDevice, copy_stream_)); } raft::resources const* res_ = nullptr; - DataT const* source_ = nullptr; - std::size_t n_rows_ = 0; std::size_t row_width_ = 0; std::size_t batch_size_ = 0; - std::size_t n_batches_ = 0; + std::vector> batches_; rmm::cuda_stream_view copy_stream_; rmm::device_uvector buffer_0_; rmm::device_uvector buffer_1_; diff --git a/cpp/src/cluster/detail/kmeans_mg.cuh b/cpp/src/cluster/detail/kmeans_mg.cuh index c6871020ca..721fc27585 100644 --- a/cpp/src/cluster/detail/kmeans_mg.cuh +++ b/cpp/src/cluster/detail/kmeans_mg.cuh @@ -6,13 +6,13 @@ #include "../kmeans.cuh" #include "kmeans.cuh" +#include "kmeans_batch_loader.cuh" #include "kmeans_common.cuh" #include "kmeans_mg_batched_init.cuh" #include "kmeans_mg_distributed_init.cuh" #include "../../core/mnmg_comms.cuh" #include "../../core/omp_wrapper.hpp" -#include "../../neighbors/detail/ann_utils.cuh" #include #include @@ -26,6 +26,7 @@ #include #include #include +#include #include #include #include @@ -37,6 +38,7 @@ #include #include +#include #include #include @@ -118,9 +120,13 @@ void mnmg_fit( { using data_part_view_t = raft::mdspan, raft::row_major, Accessor>; - using data_batch_iterator_t = - cuvs::spatial::knn::detail::utils::batch_load_iterator; constexpr bool data_on_device = raft::is_device_mdspan_v; + using input_partition_t = cuvs::cluster::kmeans::detail::kmeans_input_partition; + using data_batch_loader_t = + cuvs::cluster::kmeans::detail::kmeans_batch_loader; + using host_batch_loader_t = + cuvs::cluster::kmeans::detail::kmeans_batch_loader; + using batch_t = cuvs::cluster::kmeans::detail::kmeans_batch; bool use_nccl = raft::resource::is_multi_gpu(handle); int rank, num_ranks; @@ -208,13 +214,13 @@ void mnmg_fit( auto rank_centroids = use_nccl ? rank_centroids_arr.view() : centroids; auto new_centroids = raft::make_device_matrix(dev_res, n_clusters, n_features); auto centroid_sums = raft::make_device_matrix(dev_res, n_clusters, n_features); - auto weight_per_cluster = raft::make_device_vector(dev_res, n_clusters); - auto clustering_cost = raft::make_device_vector(dev_res, 1); - auto batch_clustering_cost = raft::make_device_vector(dev_res, 1); - auto sqrd_norm_error_dev = raft::make_device_scalar(dev_res, DataT{0}); - auto batch_cost = raft::make_device_scalar(dev_res, DataT{0}); - IndexT alloc_batch_size = device_buffer_samples; - auto batch_weights = raft::make_device_vector(dev_res, alloc_batch_size); + auto weight_per_cluster = raft::make_device_vector(dev_res, n_clusters); + auto clustering_cost = raft::make_device_vector(dev_res, 1); + auto batch_inertia = raft::make_device_scalar(dev_res, DataT{0}); + auto sqrd_norm_error_dev = raft::make_device_scalar(dev_res, DataT{0}); + auto batch_cost = raft::make_device_scalar(dev_res, DataT{0}); + IndexT alloc_batch_size = device_buffer_samples; + auto batch_weights = raft::make_device_vector(dev_res, alloc_batch_size); auto minClusterAndDistance = raft::make_device_vector, IndexT>(dev_res, alloc_batch_size); auto L2NormBatch = @@ -302,8 +308,7 @@ void mnmg_fit( auto d_prior_cost = raft::make_device_scalar(dev_res, DataT{0}); auto d_done_flag = raft::make_device_scalar(dev_res, 0); - auto h_done_flag = raft::make_host_scalar(0); - auto h_norm_cache = raft::make_host_vector(!data_on_device ? n_local : IndexT{0}); + auto h_done_flag = raft::make_pinned_scalar(dev_res, 0); auto d_norms = raft::make_device_vector(dev_res, data_on_device ? n_local : IndexT{0}); bool norms_cached = false; @@ -331,15 +336,81 @@ void mnmg_fit( } } - auto prepare_batch_weights = [&](size_t part_idx, IndexT batch_offset, IndexT cur_batch_size) - -> raft::device_vector_view { + std::vector data_inputs; + data_inputs.reserve(X_parts.size()); + for (auto const& X_part : X_parts) { + data_inputs.push_back({X_part.data_handle(), static_cast(X_part.extent(0))}); + } + + std::vector weight_inputs; + if constexpr (!data_on_device) { + if (sample_weights) { + weight_inputs.reserve(sample_weight_parts->size()); + for (auto const& weights : *sample_weight_parts) { + weight_inputs.push_back({weights.data_handle(), static_cast(weights.extent(0))}); + } + } + } + + auto batch_mr = data_on_device ? raft::resource::get_workspace_resource_ref(dev_res) + : raft::resource::get_large_workspace_resource_ref(dev_res); + auto batch_copy_stream = stream; + std::optional owned_batch_copy_stream; + if constexpr (!data_on_device) { + if (dev_res.has_resource_factory(raft::resource::resource_type::CUDA_STREAM_POOL) && + raft::resource::get_stream_pool_size(dev_res) >= 1) { + batch_copy_stream = raft::resource::get_stream_from_stream_pool(dev_res); + } else { + owned_batch_copy_stream.emplace(rmm::cuda_stream::flags::non_blocking); + batch_copy_stream = owned_batch_copy_stream->view(); + } + } + + data_batch_loader_t data_batches( + dev_res, data_inputs, n_features, device_buffer_samples, batch_copy_stream, batch_mr); + std::optional weight_batches; + if constexpr (!data_on_device) { + if (sample_weights) { + weight_batches.emplace( + dev_res, weight_inputs, IndexT{1}, device_buffer_samples, batch_copy_stream, batch_mr); + RAFT_EXPECTS(weight_batches->num_batches() == data_batches.num_batches(), + "KMeans data and weight batches do not align"); + } + } + + auto prefetch_batch = [&](std::size_t batch_pos) { + data_batches.prefetch(batch_pos); + if (weight_batches.has_value()) { 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; + }; + + auto compute_batch_norms = [&](DataT const* batch_data, IndexT batch_size) { + auto batch_view = + raft::make_device_matrix_view(batch_data, batch_size, n_features); + auto norm_view = + raft::make_device_vector_view(L2NormBatch.data_handle(), batch_size); + raft::linalg::norm( + dev_res, batch_view, norm_view); + }; + + auto prepare_batch_weights = + [&](size_t part_idx, + IndexT batch_offset, + DataT const* staged_weights, + IndexT cur_batch_size) -> raft::device_vector_view { if (sample_weights) { if constexpr (data_on_device) { return raft::make_device_vector_view( d_scaled_weights.data_handle() + part_offsets[part_idx] + batch_offset, cur_batch_size); } else { - auto const* src = (*sample_weight_parts)[part_idx].data_handle() + batch_offset; - raft::copy(batch_weights.data_handle(), src, cur_batch_size, stream); + RAFT_EXPECTS(staged_weights != nullptr, "host KMeans weights were not staged"); + raft::copy(batch_weights.data_handle(), staged_weights, cur_batch_size, stream); auto batch_weights_mut = raft::make_device_vector_view(batch_weights.data_handle(), cur_batch_size); raft::linalg::map( @@ -398,79 +469,69 @@ void mnmg_fit( raft::matrix::fill(dev_res, weight_per_cluster.view(), DataT{0}); raft::matrix::fill(dev_res, clustering_cost.view(), DataT{0}); - for (size_t part_idx = 0; part_idx < X_parts.size(); ++part_idx) { - auto const& X_part = X_parts[part_idx]; - auto part_rows = static_cast(X_part.extent(0)); - if (part_rows == 0) { continue; } + 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 weight_batch; + if (weight_batches.has_value()) { + weight_batch.emplace(weight_batches->acquire(batch_pos)); + } - data_batch_iterator_t data_batches(dev_res, - X_part, - static_cast(device_buffer_samples), - stream, - rmm::mr::get_current_device_resource_ref(), - true); - - for (auto const& data_batch : data_batches) { - IndexT current_batch_size = static_cast(data_batch.size()); - auto batch_offset = static_cast(data_batch.offset()); - - auto batch_data_view = raft::make_device_matrix_view( - data_batch.data(), current_batch_size, n_features); - - auto batch_weights_view = - prepare_batch_weights(part_idx, batch_offset, current_batch_size); - - auto norm_offset = part_offsets[part_idx] + batch_offset; - raft::device_vector_view L2NormBatch_const; - if constexpr (data_on_device) { - auto norm_slice = raft::make_device_vector_view( - d_norms.data_handle() + norm_offset, current_batch_size); - if (!norms_cached) { - raft::linalg::norm( - dev_res, batch_data_view, norm_slice); - } - L2NormBatch_const = raft::make_const_mdspan(norm_slice); - } else { - auto norm_slice = raft::make_device_vector_view( - L2NormBatch.data_handle(), current_batch_size); - if (!norms_cached) { - raft::linalg::norm( - dev_res, batch_data_view, norm_slice); - raft::copy(h_norm_cache.data_handle() + norm_offset, - L2NormBatch.data_handle(), - current_batch_size, - stream); - } else { - raft::copy(L2NormBatch.data_handle(), - h_norm_cache.data_handle() + norm_offset, - current_batch_size, - stream); - } - L2NormBatch_const = raft::make_const_mdspan(norm_slice); - } + const auto part_idx = data_batch.partition(); + const auto current_batch_size = static_cast(data_batch.size()); + const auto batch_offset = static_cast(data_batch.offset()); + const auto* staged_weights = weight_batch.has_value() ? weight_batch->data() : nullptr; - auto minClusterAndDistance_view = - raft::make_device_vector_view, IndexT>( - minClusterAndDistance.data_handle(), current_batch_size); - - cuvs::cluster::kmeans::detail::process_batch( - dev_res, - batch_data_view, - batch_weights_view, - rank_centroids_const, - metric, - params.batch_samples, - params.batch_centroids, - minClusterAndDistance_view, - L2NormBatch_const, - L2NormBuf_OR_DistBuf, - workspace, - centroid_sums.view(), - weight_per_cluster.view(), - raft::make_device_scalar_view(clustering_cost.data_handle()), - batch_workspace, - batch_cost.view()); + auto batch_data_view = raft::make_device_matrix_view( + data_batch.data(), current_batch_size, n_features); + auto batch_weights_view = + prepare_batch_weights(part_idx, batch_offset, staged_weights, current_batch_size); + + auto norm_offset = part_offsets[part_idx] + batch_offset; + raft::device_vector_view L2NormBatch_const; + if constexpr (data_on_device) { + auto norm_slice = raft::make_device_vector_view( + d_norms.data_handle() + norm_offset, current_batch_size); + if (!norms_cached) { + raft::linalg::norm( + dev_res, batch_data_view, norm_slice); + } + L2NormBatch_const = raft::make_const_mdspan(norm_slice); + } else { + compute_batch_norms(data_batch.data(), current_batch_size); + L2NormBatch_const = raft::make_device_vector_view( + L2NormBatch.data_handle(), current_batch_size); } + + // During cold fill, enqueue the first real consumer before the second H2D. Once both slots + // are active this is a no-op; recycle() keeps the copy stream one batch ahead thereafter. + prefetch_batch((batch_pos + 1) % data_batches.num_batches()); + + auto minClusterAndDistance_view = + raft::make_device_vector_view, IndexT>( + minClusterAndDistance.data_handle(), current_batch_size); + + cuvs::cluster::kmeans::detail::process_batch( + dev_res, + batch_data_view, + batch_weights_view, + rank_centroids_const, + metric, + iter_params.batch_samples, + iter_params.batch_centroids, + minClusterAndDistance_view, + L2NormBatch_const, + L2NormBuf_OR_DistBuf, + workspace, + centroid_sums.view(), + weight_per_cluster.view(), + raft::make_device_scalar_view(clustering_cost.data_handle()), + batch_workspace, + batch_cost.view()); + + 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); } } norms_cached = true; @@ -520,49 +581,92 @@ void mnmg_fit( }); raft::copy(dev_res, - h_done_flag.view(), + raft::make_pinned_scalar_view(h_done_flag.data_handle()), raft::make_device_scalar_view(d_done_flag.data_handle())); } local_n_iter = std::min(local_n_iter, static_cast(iter_params.max_iter)); raft::matrix::fill(dev_res, clustering_cost.view(), DataT{0}); - for (size_t part_idx = 0; part_idx < X_parts.size(); ++part_idx) { - auto const& X_part = X_parts[part_idx]; - auto part_rows = static_cast(X_part.extent(0)); - if (part_rows == 0) { continue; } - - data_batch_iterator_t data_batches(dev_res, - X_part, - static_cast(device_buffer_samples), - stream, - rmm::mr::get_current_device_resource_ref(), - true); - - for (auto const& data_batch : data_batches) { - IndexT current_batch_size = static_cast(data_batch.size()); - auto batch_offset = static_cast(data_batch.offset()); - - auto batch_data_view = raft::make_device_matrix_view( - data_batch.data(), current_batch_size, n_features); + 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 weight_batch; + if (weight_batches.has_value()) { weight_batch.emplace(weight_batches->acquire(batch_pos)); } + + const auto part_idx = data_batch.partition(); + const auto current_batch_size = static_cast(data_batch.size()); + const auto batch_offset = static_cast(data_batch.offset()); + const auto* staged_weights = weight_batch.has_value() ? weight_batch->data() : nullptr; + auto batch_data_view = raft::make_device_matrix_view( + data_batch.data(), current_batch_size, n_features); + + std::optional> batch_sw = std::nullopt; + if (sample_weights) { + batch_sw = + prepare_batch_weights(part_idx, batch_offset, staged_weights, current_batch_size); + } - std::optional> batch_sw = std::nullopt; - if (sample_weights) { - batch_sw = prepare_batch_weights(part_idx, batch_offset, current_batch_size); + raft::device_vector_view l2_norm_view; + if constexpr (data_on_device) { + auto norm_slice = raft::make_device_vector_view( + d_norms.data_handle() + part_offsets[part_idx] + batch_offset, current_batch_size); + if (!norms_cached) { + raft::linalg::norm( + dev_res, batch_data_view, norm_slice); } + l2_norm_view = raft::make_const_mdspan(norm_slice); + } else { + compute_batch_norms(data_batch.data(), current_batch_size); + l2_norm_view = raft::make_device_vector_view(L2NormBatch.data_handle(), + current_batch_size); + } - raft::matrix::fill(dev_res, batch_clustering_cost.view(), DataT{0}); - cuvs::cluster::kmeans::cluster_cost( - dev_res, - batch_data_view, - rank_centroids_const, - raft::make_device_scalar_view(batch_clustering_cost.data_handle()), - batch_sw); + if (batch_pos + 1 < data_batches.num_batches() || seed_iter + 1 < n_init) { + prefetch_batch((batch_pos + 1) % data_batches.num_batches()); + } - raft::linalg::add(dev_res, - raft::make_const_mdspan(clustering_cost.view()), - raft::make_const_mdspan(batch_clustering_cost.view()), - clustering_cost.view()); + auto min_cad_view = raft::make_device_vector_view, IndexT>( + minClusterAndDistance.data_handle(), current_batch_size); + cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute( + dev_res, + batch_data_view, + rank_centroids_const, + min_cad_view, + l2_norm_view, + L2NormBuf_OR_DistBuf, + cuvs::distance::DistanceType::L2Expanded, + iter_params.batch_samples, + iter_params.batch_centroids, + workspace); + if (batch_sw.has_value()) { + raft::linalg::map( + dev_res, + min_cad_view, + [] __device__(raft::KeyValuePair pair, DataT weight) { + pair.value *= weight; + return pair; + }, + raft::make_const_mdspan(min_cad_view), + batch_sw.value()); + } + cuvs::cluster::kmeans::detail::computeClusterCost( + dev_res, min_cad_view, workspace, 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); } } } comms.allreduce(clustering_cost.data_handle(), clustering_cost.data_handle(), 1);