diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index e3ffb4a439..b8fa6e8de1 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 @@ -686,25 +688,49 @@ void kmeans_fit( auto minClusterAndDistance = raft::make_device_vector, IndexT>( handle, device_buffer_samples); - auto L2NormBatch = 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); + 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 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}); } @@ -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(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( - 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); @@ -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( 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.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) { - 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); + 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)); } + IndexT cur_batch_size = static_cast(data_batch.size()); + 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); 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); } } + // 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}; 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 +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(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(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( 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) { - 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); + 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> 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_batch.has_value() ? weight_batch->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); + 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( + 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()), + 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..0891caa0d9 --- /dev/null +++ b/cpp/src/cluster/detail/kmeans_batch_loader.cuh @@ -0,0 +1,357 @@ +/* + * 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 +#include + +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 { + 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_; } + [[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, + 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; + std::size_t partition_ = 0; + std::size_t position_ = 0; + int slot_ = 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& 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) + : 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 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 < 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: + 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::vector> batches_; +}; + +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) + : 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), + row_width_(static_cast(row_width)), + 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) + { + for (std::size_t partition = 0; partition < partitions.size(); ++partition) { + append_batches(partitions[partition], partition); + } + if (batches_.empty()) { return; } + + 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 (batches_.size() > 1) { + buffer_1_.resize(row_width_ * max_batch_rows, copy_stream_); + buffer_ptrs_[1] = 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 (!batches_.empty()) { + 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 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 < 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) { + stage(slot, pos); + return; + } + } + } + + /** Make a prefetched batch visible to kernels on the main stream. */ + [[nodiscard]] auto acquire(std::size_t pos) -> kmeans_batch + { + 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; + + 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"); + } + + /** 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 < 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. + if (positions_[slot] == next_pos) { + states_[slot] = slot_state::staged; + return; + } + + 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 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 + { + 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 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.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) + { + auto const& batch = batches_[pos]; + raft::copy(dst, + batch.source + batch.offset * row_width_, + batch.size * row_width_, + copy_stream_); + } + + raft::resources const* res_ = nullptr; + std::size_t row_width_ = 0; + std::size_t batch_size_ = 0; + std::vector> batches_; + rmm::cuda_stream_view copy_stream_; + rmm::device_uvector buffer_0_; + rmm::device_uvector buffer_1_; + 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_; +}; + +} // 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..8df1d42c9c 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 @@ -118,9 +119,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,12 +213,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}); - 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 = @@ -301,8 +307,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; @@ -330,15 +335,77 @@ 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; + 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); + } + } + + 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( @@ -397,78 +464,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); + 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; @@ -518,49 +576,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); diff --git a/cpp/tests/cluster/kmeans.cu b/cpp/tests/cluster/kmeans.cu index 59051484f4..c994b9c00e 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,67 @@ 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); + + 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.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 (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); + } + } + } + + 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