diff --git a/c/include/cuvs/core/c_api.h b/c/include/cuvs/core/c_api.h index 0ed300d6ba..77b4eb1f5c 100644 --- a/c/include/cuvs/core/c_api.h +++ b/c/include/cuvs/core/c_api.h @@ -122,6 +122,25 @@ CUVS_EXPORT cuvsError_t cuvsResourcesCreateWithMemoryTracking(cuvsResources_t* r */ CUVS_EXPORT cuvsError_t cuvsResourcesDestroy(cuvsResources_t res); +/** + * @brief Set a memory pool on the device used by these resources + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] percent_of_free_memory Percentage of free device memory to allocate for the pool + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsResourcesSetMemoryPool(cuvsResources_t res, + int percent_of_free_memory); + +/** + * @brief Set a CUDA stream pool on these resources + * + * @param[in] res cuvsResources_t opaque C handle + * @param[in] num_streams Number of non-blocking CUDA streams in the pool + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsResourcesSetStreamPool(cuvsResources_t res, size_t num_streams); + /** * @brief Set cudaStream_t on cuvsResources_t to queue CUDA kernels on APIs * that accept a cuvsResources_t handle @@ -211,6 +230,16 @@ CUVS_EXPORT cuvsError_t cuvsMultiGpuResourcesDestroy(cuvsResources_t res); * @return cuvsError_t */ CUVS_EXPORT cuvsError_t cuvsMultiGpuResourcesSetMemoryPool(cuvsResources_t res, int percent_of_free_memory); + +/** + * @brief Set a CUDA stream pool on all devices managed by the multi-GPU resources + * + * @param[in] res cuvsResources_t opaque C handle for multi-GPU resources + * @param[in] num_streams Number of CUDA streams in each device's pool + * @return cuvsError_t + */ +CUVS_EXPORT cuvsError_t cuvsMultiGpuResourcesSetStreamPool(cuvsResources_t res, + size_t num_streams); /** @} */ /** diff --git a/c/src/core/c_api.cpp b/c/src/core/c_api.cpp index 27d3289e75..fd23066f0f 100644 --- a/c/src/core/c_api.cpp +++ b/c/src/core/c_api.cpp @@ -9,12 +9,16 @@ #include #include #include +#include #include #include +#include #include #include #include #include +#include +#include #include #include #include @@ -29,18 +33,78 @@ #include #include #include +#include #include #include #include +namespace { + +class single_gpu_resources : public raft::resources { + public: + ~single_gpu_resources() override { reset_memory_pool(); } + + void set_memory_pool(int percent_of_free_memory) + { + RAFT_EXPECTS(percent_of_free_memory > 0 && percent_of_free_memory <= 100, + "percent_of_free_memory must be in the range [1, 100]"); + + reset_memory_pool(); + pool_device_id_ = rmm::get_current_cuda_device(); + auto pool = rmm::mr::pool_memory_resource{ + rmm::mr::get_current_device_resource_ref(), + rmm::percent_of_free_device_memory(percent_of_free_memory)}; + previous_memory_resource_.emplace( + rmm::mr::set_per_device_resource(*pool_device_id_, std::move(pool))); + } + + private: + void reset_memory_pool() + { + if (!previous_memory_resource_.has_value()) { return; } + + rmm::cuda_set_device_raii device_guard{*pool_device_id_}; + rmm::mr::set_per_device_resource(*pool_device_id_, std::move(*previous_memory_resource_)); + previous_memory_resource_.reset(); + pool_device_id_.reset(); + } + + std::optional pool_device_id_; + std::optional previous_memory_resource_; +}; + +} // namespace + extern "C" cuvsError_t cuvsResourcesCreate(cuvsResources_t* res) { return cuvs::core::translate_exceptions([=] { - auto res_ptr = new raft::resources{}; + auto res_ptr = new single_gpu_resources{}; *res = reinterpret_cast(res_ptr); }); } +extern "C" cuvsError_t cuvsResourcesSetMemoryPool(cuvsResources_t res, + int percent_of_free_memory) +{ + return cuvs::core::translate_exceptions([=] { + auto res_ptr = dynamic_cast(reinterpret_cast(res)); + RAFT_EXPECTS(res_ptr != nullptr, + "memory pools are not supported on memory-tracking resources"); + res_ptr->set_memory_pool(percent_of_free_memory); + }); +} + +extern "C" cuvsError_t cuvsResourcesSetStreamPool(cuvsResources_t res, size_t num_streams) +{ + return cuvs::core::translate_exceptions([=] { + RAFT_EXPECTS(num_streams > 0, "num_streams must be greater than zero"); + auto res_ptr = reinterpret_cast(res); + RAFT_EXPECTS(res_ptr != nullptr, "res must not be NULL"); + raft::resource::set_cuda_stream_pool( + *res_ptr, std::make_shared(num_streams)); + }); +} + extern "C" cuvsError_t cuvsResourcesSetWorkspacePool(cuvsResources_t res, size_t initial_size_bytes) { return cuvs::core::translate_exceptions([=] { @@ -132,6 +196,24 @@ extern "C" cuvsError_t cuvsMultiGpuResourcesSetMemoryPool(cuvsResources_t res, }); } +extern "C" cuvsError_t cuvsMultiGpuResourcesSetStreamPool(cuvsResources_t res, + size_t num_streams) +{ + return cuvs::core::translate_exceptions([=] { + RAFT_EXPECTS(num_streams > 0, "num_streams must be greater than zero"); + auto res_ptr = reinterpret_cast(res); + RAFT_EXPECTS(res_ptr != nullptr, "res must not be NULL"); + + auto& device_resources = raft::resource::get_multi_gpu_resource(*res_ptr); + for (auto& device_resource : device_resources) { + rmm::cuda_set_device_raii device_guard{ + rmm::cuda_device_id{raft::resource::get_device_id(device_resource)}}; + raft::resource::set_cuda_stream_pool( + device_resource, std::make_shared(num_streams)); + } + }); +} + extern "C" cuvsError_t cuvsStreamSet(cuvsResources_t res, cudaStream_t stream) { return cuvs::core::translate_exceptions([=] { diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index e3ffb4a439..7f58beefd3 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,47 @@ 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_cost = raft::make_device_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 +782,11 @@ 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); } + }; + RAFT_LOG_DEBUG( "KMeans.fit: n_samples=%zu, n_features=%zu, n_clusters=%d, device_buffer_samples=%zu", static_cast(n_samples), @@ -767,10 +796,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 +855,45 @@ 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. + data_batches.start(); + if (weight_batches.has_value()) { weight_batches->start(); } + 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 +909,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,6 +950,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())); + // 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. } { @@ -934,29 +959,28 @@ void kmeans_fit( 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); + data_batches.start(); + if (weight_batches.has_value()) { weight_batches->start(); } + 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); } + 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, @@ -964,8 +988,18 @@ void kmeans_fit( centroids_const, raft::make_host_scalar_view(&batch_cost), batch_sw); - iter_inertia += batch_cost; + + 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); } + } } } 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..039353ef36 --- /dev/null +++ b/cpp/src/cluster/detail/kmeans_batch_loader.cuh @@ -0,0 +1,363 @@ +/* + * 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 start() noexcept {} + 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::resource::sync_stream(*res_); } + raft::resource::sync_stream(*res_, copy_stream_); + for (auto event : events_) { + if (event != nullptr) { RAFT_CUDA_TRY_NO_THROW(cudaEventDestroy(event)); } + } + } + + [[nodiscard]] auto num_batches() const noexcept -> std::size_t { return batches_.size(); } + + /** Start the pipeline by staging its first batch. */ + void start() + { + if (started_) { return; } + if (!batches_.empty()) { prefetch(0); } + started_ = true; + } + + /** 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]; + bool started_ = false; + 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..c178416080 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -683,6 +683,7 @@ __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 Reusable scratch scalar for this batch's cost */ template void process_batch( @@ -700,7 +701,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 +744,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..650a5af48d 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; @@ -211,9 +216,10 @@ void mnmg_fit( 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 batch_cost = raft::make_device_scalar(dev_res, DataT{0}); 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 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,70 @@ 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); } + }; + + 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 +457,70 @@ 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; } + data_batches.start(); + if (weight_batches.has_value()) { weight_batches->start(); } + 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 +570,59 @@ 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); - - std::optional> batch_sw = std::nullopt; - if (sample_weights) { - batch_sw = prepare_batch_weights(part_idx, batch_offset, current_batch_size); - } + data_batches.start(); + if (weight_batches.has_value()) { weight_batches->start(); } + 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); + } - 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()); + 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); + + raft::linalg::add(dev_res, + raft::make_const_mdspan(clustering_cost.view()), + raft::make_const_mdspan(batch_clustering_cost.view()), + clustering_cost.view()); + + 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..1099e14348 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,69 @@ 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.start(); + // Starting an active pipeline is a no-op. + loader.start(); + 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 diff --git a/python/cuvs/cuvs/common/c_api.pxd b/python/cuvs/cuvs/common/c_api.pxd index 02e7f7a289..e197385525 100644 --- a/python/cuvs/cuvs/common/c_api.pxd +++ b/python/cuvs/cuvs/common/c_api.pxd @@ -6,6 +6,7 @@ from cuda.bindings.cyruntime cimport cudaStream_t +from libc.stddef cimport size_t from libc.stdint cimport int64_t, uintptr_t from cuvs.common.cydlpack cimport DLManagedTensor @@ -24,6 +25,10 @@ cdef extern from "cuvs/core/c_api.h": const char* csv_path, int64_t sample_interval_ms) cuvsError_t cuvsResourcesDestroy(cuvsResources_t res) + cuvsError_t cuvsResourcesSetMemoryPool(cuvsResources_t res, + int percent_of_free_memory) + cuvsError_t cuvsResourcesSetStreamPool(cuvsResources_t res, + size_t num_streams) cuvsError_t cuvsStreamSet(cuvsResources_t res, cudaStream_t stream) cuvsError_t cuvsStreamSync(cuvsResources_t res) const char * cuvsGetLastErrorText() @@ -35,6 +40,8 @@ cdef extern from "cuvs/core/c_api.h": cuvsError_t cuvsMultiGpuResourcesDestroy(cuvsResources_t res) cuvsError_t cuvsMultiGpuResourcesSetMemoryPool(cuvsResources_t res, int percent_of_free_memory) + cuvsError_t cuvsMultiGpuResourcesSetStreamPool(cuvsResources_t res, + size_t num_streams) cuvsError_t cuvsMatrixCopy(cuvsResources_t res, DLManagedTensor * src, DLManagedTensor * dst) diff --git a/python/cuvs/cuvs/common/mg_resources.pyx b/python/cuvs/cuvs/common/mg_resources.pyx index 4b9e65df61..dc33318b6b 100644 --- a/python/cuvs/cuvs/common/mg_resources.pyx +++ b/python/cuvs/cuvs/common/mg_resources.pyx @@ -11,6 +11,7 @@ from cuvs.common.c_api cimport ( cuvsMultiGpuResourcesCreateWithDeviceIds, cuvsMultiGpuResourcesDestroy, cuvsMultiGpuResourcesSetMemoryPool, + cuvsMultiGpuResourcesSetStreamPool, cuvsResources_t, cuvsStreamSet, cuvsStreamSync, @@ -108,6 +109,20 @@ cdef class MultiGpuResources: check_cuvs(cuvsMultiGpuResourcesSetMemoryPool( self.c_obj, percent_of_free_memory)) + def set_stream_pool(self, num_streams=1): + """ + Set a CUDA stream pool on all devices managed by these resources. + + Parameters + ---------- + num_streams : int, default=1 + Number of non-blocking CUDA streams in each device's pool. + """ + if num_streams <= 0: + raise ValueError("num_streams must be greater than zero") + check_cuvs(cuvsMultiGpuResourcesSetStreamPool( + self.c_obj, num_streams)) + def get_c_obj(self): """ Return the pointer to the underlying c_obj as a size_t diff --git a/python/cuvs/cuvs/common/resources.pyx b/python/cuvs/cuvs/common/resources.pyx index 00976d9ad8..5e451f7530 100644 --- a/python/cuvs/cuvs/common/resources.pyx +++ b/python/cuvs/cuvs/common/resources.pyx @@ -14,6 +14,8 @@ from cuvs.common.c_api cimport ( cuvsResourcesCreate, cuvsResourcesCreateWithMemoryTracking, cuvsResourcesDestroy, + cuvsResourcesSetMemoryPool, + cuvsResourcesSetStreamPool, cuvsStreamSet, cuvsStreamSync, ) @@ -89,6 +91,31 @@ cdef class Resources: def sync(self): check_cuvs(cuvsStreamSync(self.c_obj)) + def set_memory_pool(self, percent_of_free_memory): + """ + Set a memory pool on the device used by these resources. + + Parameters + ---------- + percent_of_free_memory : int + Percentage of free device memory to allocate for the pool. + """ + check_cuvs(cuvsResourcesSetMemoryPool( + self.c_obj, percent_of_free_memory)) + + def set_stream_pool(self, num_streams=1): + """ + Set a CUDA stream pool on these resources. + + Parameters + ---------- + num_streams : int, default=1 + Number of non-blocking CUDA streams in the pool. + """ + if num_streams <= 0: + raise ValueError("num_streams must be greater than zero") + check_cuvs(cuvsResourcesSetStreamPool(self.c_obj, num_streams)) + def get_c_obj(self): """ Return the pointer to the underlying c_obj as a size_t