diff --git a/cpp/src/cluster/detail/minClusterDistanceCompute.cu b/cpp/src/cluster/detail/minClusterDistanceCompute.cu index b15119599e..ee3cc3cdfd 100644 --- a/cpp/src/cluster/detail/minClusterDistanceCompute.cu +++ b/cpp/src/cluster/detail/minClusterDistanceCompute.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -27,15 +27,15 @@ void minClusterAndDistanceCompute( int batch_centroids, rmm::device_uvector& workspace) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle); - auto n_samples = X.extent(0); - auto n_features = X.extent(1); - auto n_clusters = centroids.extent(0); - bool is_fused = metric == cuvs::distance::DistanceType::L2Expanded || - metric == cuvs::distance::DistanceType::L2SqrtExpanded || - metric == cuvs::distance::DistanceType::CosineExpanded; - - if (is_fused) { + cudaStream_t stream = raft::resource::get_cuda_stream(handle); + auto n_samples = X.extent(0); + auto n_features = X.extent(1); + auto n_clusters = centroids.extent(0); + const bool can_use_fused_path = metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::CosineExpanded; + + if (can_use_fused_path) { L2NormBuf_OR_DistBuf.resize(n_clusters, stream); auto centroidsNorm = raft::make_device_vector_view(L2NormBuf_OR_DistBuf.data(), n_clusters); @@ -51,10 +51,10 @@ void minClusterAndDistanceCompute( raft::KeyValuePair initial_value(0, std::numeric_limits::max()); raft::matrix::fill(handle, minClusterAndDistance, initial_value); - bool should_use_fused = + const bool use_fused_path = use_fused(handle, n_samples, n_clusters, n_features); - if (should_use_fused) { + if (use_fused_path) { workspace.resize((sizeof(int)) * n_samples, stream); cuvs::distance::fusedDistanceNNMinReduce, IndexT>( @@ -74,26 +74,66 @@ void minClusterAndDistanceCompute( 0.0f, stream); } else { - workspace.resize(sizeof(DataT) * n_samples * n_clusters, stream); - - cuvs::distance:: - unfusedDistanceNNMinReduce, IndexT>( - handle, - minClusterAndDistance.data_handle(), - X.data_handle(), - centroids.data_handle(), - L2NormX.data_handle(), - centroidsNorm.data_handle(), - n_samples, - n_clusters, - n_features, - (void*)workspace.data(), - metric != cuvs::distance::DistanceType::L2Expanded, - false, - true, - metric, - 0.0f, - stream); + auto dataBatchSize = getDataBatchSize(batch_samples, n_samples); + auto centroidsBatchSize = getCentroidsBatchSize(batch_centroids, n_clusters); + + // The unfused reduction indexes its distance matrix with IndexT. + dataBatchSize = + std::min(dataBatchSize, std::numeric_limits::max() / centroidsBatchSize); + + workspace.resize(sizeof(DataT) * dataBatchSize * centroidsBatchSize, stream); + + using KeyValueT = raft::KeyValuePair; + const bool tileCentroids = centroidsBatchSize < n_clusters; + rmm::device_uvector batchMinClusterAndDistance(tileCentroids ? dataBatchSize : 0, + stream); + + for (IndexT dIdx = 0; dIdx < n_samples;) { + auto ns = std::min(dataBatchSize, n_samples - dIdx); + auto minClusterAndDistanceView = raft::make_device_vector_view( + minClusterAndDistance.data_handle() + dIdx, ns); + + for (IndexT cIdx = 0; cIdx < n_clusters;) { + auto nc = std::min(centroidsBatchSize, n_clusters - cIdx); + auto batchMin = tileCentroids ? batchMinClusterAndDistance.data() + : minClusterAndDistanceView.data_handle(); + + cuvs::distance::unfusedDistanceNNMinReduce( + handle, + batchMin, + X.data_handle() + dIdx * n_features, + centroids.data_handle() + cIdx * n_features, + L2NormX.data_handle() + dIdx, + centroidsNorm.data_handle() + cIdx, + ns, + nc, + n_features, + (void*)workspace.data(), + metric != cuvs::distance::DistanceType::L2Expanded, + tileCentroids, + true, + metric, + 0.0f, + stream); + + if (tileCentroids) { + // Convert tile-local centroid indices and merge the tile minima. + auto batchMinView = + raft::make_device_vector_view(batchMin, ns); + raft::linalg::map( + handle, + minClusterAndDistanceView, + [cIdx] __device__(KeyValueT current, KeyValueT batch) { + batch.key += cIdx; + return batch.value < current.value ? batch : current; + }, + raft::make_const_mdspan(minClusterAndDistanceView), + batchMinView); + } + cIdx += nc; + } + dIdx += ns; + } } } else { auto dataBatchSize = getDataBatchSize(batch_samples, n_samples);