diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index e3ffb4a439..39b638d762 100644 --- a/cpp/src/cluster/detail/kmeans.cuh +++ b/cpp/src/cluster/detail/kmeans.cuh @@ -24,6 +24,7 @@ #include #include #include +#include #include #include #include @@ -39,6 +40,7 @@ #include #include +#include #include #include @@ -686,25 +688,58 @@ 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 batch_weights_buf = raft::make_device_vector(handle, device_buffer_samples); + auto minClusterDistance = raft::make_device_vector(handle, device_buffer_samples); + auto L2NormBatch = raft::make_device_vector(handle, device_buffer_samples); + 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}); 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_memory = raft::resource::get_workspace_resource_ref(handle); + if constexpr (!data_on_device) { + size_t batch_staging_bytes = + static_cast(device_buffer_samples) * static_cast(n_features) * sizeof(DataT); + if (weight_ptr != nullptr) { + batch_staging_bytes += static_cast(device_buffer_samples) * sizeof(DataT); + } + if (batch_staging_bytes > raft::resource::get_workspace_free_bytes(handle)) { + batch_memory = raft::resource::get_large_workspace_resource_ref(handle); + } + } + + // Use the caller's stream pool for host-input staging. If no pool is configured, this falls + // back to the main stream and disables cross-stream prefetching. + auto [batch_copy_stream, enable_batch_prefetch] = + cuvs::spatial::knn::detail::utils::get_prefetch_stream(handle); + + auto data_batches = + cuvs::spatial::knn::detail::utils::make_batch_load_iterator(handle, + X.data_handle(), + n_samples, + n_features, + device_buffer_samples, + batch_copy_stream, + batch_memory, + enable_batch_prefetch); // Host-path weight batches: only materialized when weights are provided and // the data resides on host 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 = + cuvs::spatial::knn::detail::utils::make_batch_load_iterator(handle, + weight_ptr, + n_samples, + IndexT{1}, + device_buffer_samples, + batch_copy_stream, + batch_memory, + enable_batch_prefetch); } else { raft::matrix::fill(handle, batch_weights_buf.view(), DataT{1}); } @@ -836,19 +871,18 @@ void kmeans_fit( raft::make_device_matrix_view(new_centroids_ptr, n_clusters, n_features); data_batches.reset(); + data_batches.prefetch_next_batch(); 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(); + wt_it->prefetch_next_batch(); } 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); - } + if (wt_it.has_value()) { wt_data = (**wt_it).data(); } auto batch_data_view = raft::make_device_matrix_view( data_batch.data(), cur_batch_size, n_features); @@ -893,6 +927,11 @@ void kmeans_fit( weight_per_cluster.view(), clustering_cost.view(), batch_workspace); + data_batches.prefetch_next_batch(); + if (wt_it.has_value()) { + wt_it->prefetch_next_batch(); + ++(*wt_it); + } } if (need_compute_norms) { norms_cached = true; } @@ -933,40 +972,73 @@ void kmeans_fit( auto centroids_const = raft::make_device_matrix_view( cur_centroids_ptr, n_clusters, n_features); - iter_inertia = DataT{0}; + DataT zero = DataT{0}; + raft::copy(clustering_cost.data_handle(), &zero, 1, stream); data_batches.reset(); + data_batches.prefetch_next_batch(); 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(); + wt_it->prefetch_next_batch(); } 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); - } + if (wt_it.has_value()) { wt_data = (**wt_it).data(); } auto batch_data_view = raft::make_device_matrix_view( data_batch.data(), cur_batch_size, n_features); - std::optional> batch_sw = std::nullopt; + if constexpr (!data_on_device) { + if (need_compute_norms && norms_cached) { + raft::copy(L2NormBatch.data_handle(), + h_norm_cache.data_handle() + data_batch.offset(), + cur_batch_size, + stream); + } else { + compute_batch_norms(data_batch.data(), cur_batch_size); + } + } else { + compute_batch_norms(data_batch.data(), cur_batch_size); + } + 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); + std::optional> batch_sample_weight = + std::nullopt; if (weight_ptr != nullptr) { - batch_sw = + batch_sample_weight = cur_batch_weights(static_cast(data_batch.offset()), wt_data, cur_batch_size); } - 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; + cluster_cost(handle, + batch_data_view, + centroids_const, + min_distance_view, + l2_norm_view, + L2NormBuf_OR_DistBuf, + cuvs::distance::DistanceType::L2Expanded, + cur_batch_size, + n_clusters, + ws, + batch_inertia.view(), + batch_sample_weight); + raft::linalg::add(clustering_cost.data_handle(), + clustering_cost.data_handle(), + batch_inertia.data_handle(), + 1, + stream); + data_batches.prefetch_next_batch(); + if (wt_it.has_value()) { + wt_it->prefetch_next_batch(); + ++(*wt_it); + } } + raft::copy(&iter_inertia, clustering_cost.data_handle(), 1, stream); + raft::resource::sync_stream(handle); } if (iter_inertia < inertia[0]) { diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index ab3ef0a05a..09f8e90c67 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, diff --git a/cpp/src/cluster/kmeans.cuh b/cpp/src/cluster/kmeans.cuh index f6e2c7d819..11c0d49269 100644 --- a/cpp/src/cluster/kmeans.cuh +++ b/cpp/src/cluster/kmeans.cuh @@ -343,7 +343,6 @@ void cluster_cost( auto stream = raft::resource::get_cuda_stream(handle); auto n_clusters = centroids.extent(0); auto n_samples = X.extent(0); - auto n_features = X.extent(1); rmm::device_uvector workspace(n_samples * sizeof(IndexT), stream); @@ -353,31 +352,18 @@ void cluster_cost( auto min_cluster_distance = raft::make_device_vector(handle, n_samples); rmm::device_uvector l2_norm_or_distance_buffer(0, stream); - auto metric = cuvs::distance::DistanceType::L2Expanded; - - cuvs::cluster::kmeans::min_cluster_distance( - handle, - X, - raft::make_device_matrix_view( - const_cast(centroids.data_handle()), n_clusters, n_features), - min_cluster_distance.view(), - x_norms.view(), - l2_norm_or_distance_buffer, - metric, - n_samples, - n_clusters, - workspace); - - if (sample_weight.has_value()) { - raft::linalg::map(handle, - min_cluster_distance.view(), - raft::mul_op{}, - raft::make_const_mdspan(min_cluster_distance.view()), - sample_weight.value()); - } - - cuvs::cluster::kmeans::cluster_cost( - handle, min_cluster_distance.view(), workspace, cost, raft::add_op{}); + cuvs::cluster::kmeans::detail::cluster_cost(handle, + X, + centroids, + min_cluster_distance.view(), + x_norms.view(), + l2_norm_or_distance_buffer, + cuvs::distance::DistanceType::L2Expanded, + n_samples, + n_clusters, + workspace, + cost, + sample_weight); } /**