diff --git a/cpp/src/stats/detail/batched/silhouette_score.cuh b/cpp/src/stats/detail/batched/silhouette_score.cuh index d438baeb66..bb8a75dbf9 100644 --- a/cpp/src/stats/detail/batched/silhouette_score.cuh +++ b/cpp/src/stats/detail/batched/silhouette_score.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -128,9 +128,9 @@ rmm::device_uvector get_pairwise_distance(raft::resources const& handle value_idx& n_left_rows, value_idx& n_right_rows, value_idx& n_cols, - cuvs::distance::DistanceType metric, - cudaStream_t stream) + cuvs::distance::DistanceType metric) { + auto stream = raft::resource::get_cuda_stream(handle); rmm::device_uvector distances(n_left_rows * n_right_rows, stream); cuvs::distance::pairwise_distance( @@ -218,6 +218,8 @@ value_t silhouette_score( ++n_iters; auto chunk_stream = raft::resource::get_next_usable_stream(handle, i + chunk * j); + raft::resources chunk_handle(handle); + raft::resource::set_cuda_stream(chunk_handle, chunk_stream); const auto* left_begin = X + (i * n_cols); const auto* right_begin = X + (j * n_cols); @@ -226,7 +228,7 @@ value_t silhouette_score( auto n_right_rows = (j + chunk) < n_rows ? chunk : (n_rows - j); rmm::device_uvector distances = get_pairwise_distance( - handle, left_begin, right_begin, n_left_rows, n_right_rows, n_cols, metric, chunk_stream); + chunk_handle, left_begin, right_begin, n_left_rows, n_right_rows, n_cols, metric); compute_chunked_a_b(handle, a_ptr, diff --git a/cpp/tests/stats/silhouette_score.cu b/cpp/tests/stats/silhouette_score.cu index 279298b263..2f0d35450c 100644 --- a/cpp/tests/stats/silhouette_score.cu +++ b/cpp/tests/stats/silhouette_score.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #include "../test_utils.cuh" @@ -7,14 +7,18 @@ #include #include #include +#include #include +#include #include #include #include +#include #include +#include #include namespace cuvs { @@ -220,5 +224,42 @@ TEST_P(silhouetteScoreTestClass, Result) } INSTANTIATE_TEST_CASE_P(silhouetteScore, silhouetteScoreTestClass, ::testing::ValuesIn(inputs)); +TEST(silhouetteScore, BatchedStreamPoolOrdering) +{ + constexpr int64_t n_rows = 4096; + constexpr int64_t n_cols = 2; + constexpr int n_labels = 2; + + std::vector X(n_rows * n_cols); + std::vector labels(n_rows); + for (int64_t i = 0; i < n_rows; ++i) { + X[2 * i] = std::sin(0.01f * i); + X[2 * i + 1] = std::cos(0.013f * i); + labels[i] = i % n_labels; + } + + raft::resources handle; + raft::resource::set_cuda_stream_pool(handle, std::make_shared(4)); + auto stream = raft::resource::get_cuda_stream(handle); + + rmm::device_uvector d_X(X.size(), stream); + rmm::device_uvector d_labels(labels.size(), stream); + raft::update_device(d_X.data(), X.data(), X.size(), stream); + raft::update_device(d_labels.data(), labels.data(), labels.size(), stream); + + auto X_view = raft::make_device_matrix_view(d_X.data(), n_rows, n_cols); + auto labels_view = raft::make_device_vector_view(d_labels.data(), n_rows); + constexpr auto metric = cuvs::distance::DistanceType::L2SqrtUnexpanded; + + auto expected = + cuvs::stats::silhouette_score(handle, X_view, labels_view, std::nullopt, n_labels, metric); + + for (int repeat = 0; repeat < 8; ++repeat) { + auto actual = cuvs::stats::silhouette_score_batched( + handle, X_view, labels_view, std::nullopt, n_labels, n_rows, metric); + ASSERT_NEAR(actual, expected, 1e-4f); + } +} + } // end namespace stats } // end namespace cuvs