diff --git a/cpp/src/cluster/detail/kmeans.cuh b/cpp/src/cluster/detail/kmeans.cuh index 332fed16b5..4478cec2a6 100644 --- a/cpp/src/cluster/detail/kmeans.cuh +++ b/cpp/src/cluster/detail/kmeans.cuh @@ -1082,7 +1082,8 @@ void kmeans_fit( compute_centroid_shift(handle, raft::make_const_mdspan(centroids_const), raft::make_const_mdspan(new_centroids_view), - sqrdNormError.view()); + sqrdNormError.view(), + ws); std::swap(cur_centroids_ptr, new_centroids_ptr); diff --git a/cpp/src/cluster/detail/kmeans_centroid_sums.cuh b/cpp/src/cluster/detail/kmeans_centroid_sums.cuh new file mode 100644 index 0000000000..6248e00ae1 --- /dev/null +++ b/cpp/src/cluster/detail/kmeans_centroid_sums.cuh @@ -0,0 +1,416 @@ +/* + * 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 + +/** + * Deterministic per-cluster (weighted) row sums for the k-means centroid updates. + * + * Accumulating the rows of each cluster with atomics (as raft::linalg::reduce_rows_by_key does) + * makes the floating-point summation order, and therefore the centroids, vary from run to run. + * Through the convergence check this can change the number of Lloyd iterations, so two identical + * fits can return noticeably different centroids. Here the summation order is a function of the + * input only: + * + * 1. Row indices are stably sorted by label (cub::DeviceRadixSort), so each cluster becomes a + * contiguous segment. + * 2. Each segment is cut into pieces of at most `kRowsPerPiece` rows. One thread block reduces a + * piece with a fixed assignment of rows to threads and a fixed-order combine. + * 3. One thread block per cluster adds up that cluster's piece sums in the same fixed way, and + * adds the result to the output. + * + * The order depends only on the labels and the sizes, not on scheduling or on the launch + * configuration. + * + * Why not a CCCL reduction: cub::DeviceSegmentedReduce (run_to_run) and + * cub::DeviceReduce::ReduceByKey (run_to_run) reduce one scalar per segment or run. A + * [n_rows x n_cols] matrix would need n_clusters * n_cols segments that read the gathered rows + * column by column (strided, uncoalesced), and DeviceSegmentedReduce gives each segment to one + * thread block, which serializes large clusters when n_clusters is small. The per-cluster weight + * sum alone could use them, but it comes for free here in the same pass over the rows. + */ +namespace cuvs::cluster::kmeans::detail::sums_by_label { + +constexpr int kBlockThreads = 256; +constexpr int kRowsPerPiece = 256; + +/** Number of threads along the columns: smallest power of two >= n_cols (at most kBlockThreads). */ +inline int column_threads(int64_t n_cols) +{ + int tc = 1; + while (tc < n_cols && tc < kBlockThreads) { + tc <<= 1; + } + return tc; +} + +/** + * Sums `values(i, j)` over i in [begin, end) for each column j in [0, n_cols), using a fixed + * assignment of items to threads and a fixed-order combine, so the result does not depend on + * scheduling. The result for column j is passed to `write(j, sum)` by a single thread. + */ +template +__device__ void fixed_order_column_sums( + IndexT begin, IndexT end, IndexT n_cols, int col_threads, ValueOp values, WriteOp write) +{ + __shared__ DataT smem[kBlockThreads]; + const int lanes = kBlockThreads / col_threads; + const int tc = threadIdx.x % col_threads; + const int lane = threadIdx.x / col_threads; + for (IndexT col0 = 0; col0 < n_cols; col0 += col_threads) { + const IndexT j = col0 + tc; + DataT acc = DataT{0}; + if (j < n_cols) { +#pragma unroll 4 + for (IndexT i = begin + lane; i < end; i += lanes) { + acc += values(i, j); + } + } + smem[threadIdx.x] = acc; + __syncthreads(); + if (lane == 0 && j < n_cols) { + DataT sum = smem[tc]; + for (int l = 1; l < lanes; ++l) { + sum += smem[l * col_threads + tc]; + } + write(j, sum); + } + __syncthreads(); + } +} + +template +RAFT_KERNEL init_sort_kernel(LabelsIterator labels, IndexT n_rows, uint32_t* keys, IndexT* rows) +{ + for (IndexT i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; i < n_rows; + i += static_cast(blockDim.x) * gridDim.x) { + keys[i] = static_cast(labels[i]); + rows[i] = i; + } +} + +/** Finds each cluster's segment of the sorted labels and the number of pieces it is cut into. */ +template +RAFT_KERNEL segment_kernel(const uint32_t* sorted_keys, + IndexT n_rows, + IndexT n_clusters, + IndexT* cluster_offsets, + IndexT* piece_counts) +{ + auto lower_bound = [=](IndexT c) { + IndexT lo = 0, hi = n_rows; + while (lo < hi) { + IndexT mid = lo + (hi - lo) / 2; + if (static_cast(sorted_keys[mid]) < c) { + lo = mid + 1; + } else { + hi = mid; + } + } + return lo; + }; + for (IndexT c = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; c <= n_clusters; + c += static_cast(blockDim.x) * gridDim.x) { + IndexT begin = lower_bound(c); + cluster_offsets[c] = begin; + if (c < n_clusters) { + IndexT end = lower_bound(c + 1); + piece_counts[c] = raft::ceildiv(end - begin, kRowsPerPiece); + } else { + piece_counts[c] = 0; + } + } +} + +template +RAFT_KERNEL piece_sums_kernel(DataIterator X, + IndexT n_cols, + const DataT* weights, + bool sum_weights, + const IndexT* sorted_rows, + const IndexT* cluster_offsets, + const IndexT* piece_offsets, + IndexT n_clusters, + int col_threads, + DataT* piece_sums) +{ + const IndexT n_pieces = piece_offsets[n_clusters]; + for (IndexT piece = blockIdx.x; piece < n_pieces; piece += gridDim.x) { + // The cluster of this piece is the last c with piece_offsets[c] <= piece. + IndexT lo = 0, hi = n_clusters; + while (hi - lo > 1) { + IndexT mid = lo + (hi - lo) / 2; + if (piece_offsets[mid] <= piece) { + lo = mid; + } else { + hi = mid; + } + } + const IndexT c = lo; + const IndexT begin = cluster_offsets[c] + (piece - piece_offsets[c]) * kRowsPerPiece; + const IndexT end = std::min(begin + kRowsPerPiece, cluster_offsets[c + 1]); + DataT* out = piece_sums + static_cast(piece) * (n_cols + 1); + fixed_order_column_sums( + begin, + end, + n_cols, + col_threads, + [=](IndexT i, IndexT j) { + const IndexT row = sorted_rows[i]; + const auto x = static_cast(X[static_cast(row) * n_cols + j]); + return weights != nullptr ? x * weights[row] : x; + }, + [=](IndexT j, DataT sum) { out[j] = sum; }); + // The sum of weights is stored after the column sums. + if (sum_weights) { + fixed_order_column_sums( + begin, + end, + IndexT{1}, + 1, + [=](IndexT i, IndexT) { return weights != nullptr ? weights[sorted_rows[i]] : DataT{1}; }, + [=](IndexT, DataT sum) { out[n_cols] = sum; }); + } + } +} + +template +RAFT_KERNEL cluster_sums_kernel(const DataT* piece_sums, + IndexT n_cols, + const IndexT* piece_offsets, + IndexT n_clusters, + int col_threads, + bool reset_sums, + DataT* centroid_sums, + DataT* weight_per_cluster) +{ + for (IndexT c = blockIdx.x; c < n_clusters; c += gridDim.x) { + const IndexT p_begin = piece_offsets[c]; + const IndexT p_end = piece_offsets[c + 1]; + const IndexT stride = n_cols + 1; + DataT* out = centroid_sums + static_cast(c) * n_cols; + fixed_order_column_sums( + p_begin, + p_end, + n_cols, + col_threads, + [=](IndexT p, IndexT j) { return piece_sums[static_cast(p) * stride + j]; }, + [=](IndexT j, DataT sum) { out[j] = reset_sums ? sum : out[j] + sum; }); + if (weight_per_cluster != nullptr) { + fixed_order_column_sums( + p_begin, + p_end, + IndexT{1}, + 1, + [=](IndexT p, IndexT) { return piece_sums[static_cast(p) * stride + n_cols]; }, + [=](IndexT, DataT sum) { + weight_per_cluster[c] = reset_sums ? sum : weight_per_cluster[c] + sum; + }); + } + } +} + +inline size_t align_bytes(size_t bytes) { return raft::alignTo(bytes, 256); } + +/** Rows per pass of `weighted_sums_by_label`; bounds the workspace and fits 32-bit indices. */ +constexpr int64_t kMaxRowsPerPass = int64_t{1} << 22; + +/** One pass of `weighted_sums_by_label` over at most kMaxRowsPerPass rows. */ +template +void weighted_sums_by_label_pass(raft::resources const& handle, + DataIterator X, + const DataT* weights, + LabelsIterator labels, + uint32_t n_rows, + uint32_t n_cols, + uint32_t n_clusters, + DataT* centroid_sums, + DataT* weight_per_cluster, + rmm::device_uvector& workspace, + bool reset_sums) +{ + using IdxT = uint32_t; + cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); + + int end_bit = 1; + while (end_bit < 32 && (uint64_t{1} << end_bit) < static_cast(n_clusters)) { + ++end_bit; + } + const IdxT max_pieces = raft::ceildiv(n_rows, kRowsPerPiece) + n_clusters; + + // Workspace layout + size_t sort_bytes = 0; + RAFT_CUDA_TRY(cub::DeviceRadixSort::SortPairs(nullptr, + sort_bytes, + static_cast(nullptr), + static_cast(nullptr), + static_cast(nullptr), + static_cast(nullptr), + n_rows, + 0, + end_bit, + stream)); + size_t scan_bytes = 0; + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum(nullptr, + scan_bytes, + static_cast(nullptr), + static_cast(nullptr), + n_clusters + 1, + stream)); + const size_t keys_bytes = align_bytes(sizeof(uint32_t) * n_rows); + const size_t rows_bytes = align_bytes(sizeof(IdxT) * n_rows); + const size_t offsets_bytes = align_bytes(sizeof(IdxT) * (n_clusters + 1)); + const size_t pieces_bytes = + align_bytes(sizeof(DataT) * static_cast(max_pieces) * (n_cols + 1)); + const size_t temp_bytes = align_bytes(std::max(sort_bytes, scan_bytes)); + const size_t total_bytes = + 2 * keys_bytes + 2 * rows_bytes + 3 * offsets_bytes + pieces_bytes + temp_bytes; + if (workspace.size() < total_bytes) { workspace.resize(total_bytes, stream); } + + char* ptr = workspace.data(); + auto take = [&ptr](size_t bytes) { + char* p = ptr; + ptr += bytes; + return p; + }; + auto* keys_in = reinterpret_cast(take(keys_bytes)); + auto* keys_out = reinterpret_cast(take(keys_bytes)); + auto* rows_in = reinterpret_cast(take(rows_bytes)); + auto* rows_out = reinterpret_cast(take(rows_bytes)); + auto* cluster_offsets = reinterpret_cast(take(offsets_bytes)); + auto* piece_counts = reinterpret_cast(take(offsets_bytes)); + auto* piece_offsets = reinterpret_cast(take(offsets_bytes)); + auto* piece_sums = reinterpret_cast(take(pieces_bytes)); + void* temp = take(temp_bytes); + + if (n_rows > 0) { + const int grid = static_cast( + std::min(raft::ceildiv(n_rows, kBlockThreads), int64_t{65536})); + init_sort_kernel<<>>(labels, n_rows, keys_in, rows_in); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + RAFT_CUDA_TRY(cub::DeviceRadixSort::SortPairs( + temp, sort_bytes, keys_in, keys_out, rows_in, rows_out, n_rows, 0, end_bit, stream)); + } + { + const int grid = static_cast( + std::min(raft::ceildiv(n_clusters + 1, kBlockThreads), int64_t{65536})); + segment_kernel<<>>( + keys_out, n_rows, n_clusters, cluster_offsets, piece_counts); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + RAFT_CUDA_TRY(cub::DeviceScan::ExclusiveSum( + temp, scan_bytes, piece_counts, piece_offsets, n_clusters + 1, stream)); + + const int col_threads = column_threads(n_cols); + { + const int grid = static_cast(std::min(max_pieces, int64_t{65536})); + piece_sums_kernel<<>>(X, + n_cols, + weights, + weight_per_cluster != nullptr, + rows_out, + cluster_offsets, + piece_offsets, + n_clusters, + col_threads, + piece_sums); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } + { + const int grid = static_cast(std::min(n_clusters, int64_t{65536})); + cluster_sums_kernel<<>>(piece_sums, + n_cols, + piece_offsets, + n_clusters, + col_threads, + reset_sums, + centroid_sums, + weight_per_cluster); + RAFT_CUDA_TRY(cudaPeekAtLastError()); + } +} + +/** + * @brief Deterministic (weighted) sum of the rows of X per label, and of the weights per label. + * + * centroid_sums[c, :] (+)= sum_{i : labels[i] == c} weights[i] * X[i, :] + * weight_per_cluster[c] (+)= sum_{i : labels[i] == c} weights[i] + * + * The summation order depends only on the labels and the sizes, not on scheduling, + * so the result is bitwise reproducible. Rows are processed in passes of at most kMaxRowsPerPass + * rows, which bounds the workspace (16 bytes per row of a pass, plus the piece sums). + * + * @tparam DataT Accumulation and output type + * @tparam IndexT Index type + * @tparam DataIterator Random-access iterator over the row-major input; its values are converted + * to DataT (e.g. a transform iterator that maps int8 data to float) + * @tparam LabelsIterator Random-access iterator over the labels + * + * @param[in] handle RAFT resources handle + * @param[in] X Row-major input [n_rows x n_cols] + * @param[in] weights Per-row weights [n_rows], or nullptr for unit weights + * @param[in] labels Label of each row, in [0, n_clusters) + * @param[in] n_rows Number of rows of X + * @param[in] n_cols Number of columns of X + * @param[in] n_clusters Number of labels + * @param[inout] centroid_sums Row-major output [n_clusters x n_cols] + * @param[inout] weight_per_cluster Output [n_clusters], or nullptr to skip the weight sums + * @param[inout] workspace Scratch buffer, resized as needed + * @param[in] reset_sums Overwrite the outputs if true, add to them otherwise + */ +template +void weighted_sums_by_label(raft::resources const& handle, + DataIterator X, + const DataT* weights, + LabelsIterator labels, + IndexT n_rows, + IndexT n_cols, + IndexT n_clusters, + DataT* centroid_sums, + DataT* weight_per_cluster, + rmm::device_uvector& workspace, + bool reset_sums) +{ + if (n_clusters <= 0) { return; } + RAFT_EXPECTS(static_cast(n_clusters) < (uint64_t{1} << 31) && + static_cast(n_cols) < (uint64_t{1} << 31), + "weighted_sums_by_label: too many clusters or columns"); + // Passes are accumulated in a fixed order, so the result stays deterministic. There is always at + // least one pass, so that the outputs are reset even when there are no rows. + const auto total_rows = static_cast(n_rows); + int64_t start = 0; + do { + const int64_t pass_rows = std::min(kMaxRowsPerPass, total_rows - start); + weighted_sums_by_label_pass(handle, + X + start * static_cast(n_cols), + weights != nullptr ? weights + start : nullptr, + labels + start, + static_cast(pass_rows), + static_cast(n_cols), + static_cast(n_clusters), + centroid_sums, + weight_per_cluster, + workspace, + reset_sums && start == 0); + start += pass_rows; + } while (start < total_rows); +} + +} // namespace cuvs::cluster::kmeans::detail::sums_by_label diff --git a/cpp/src/cluster/detail/kmeans_common.cuh b/cpp/src/cluster/detail/kmeans_common.cuh index 78edf8759e..efedd05b1b 100644 --- a/cpp/src/cluster/detail/kmeans_common.cuh +++ b/cpp/src/cluster/detail/kmeans_common.cuh @@ -6,6 +6,7 @@ #include "../../distance/distance.cuh" #include "../../distance/top_1_nn.cuh" +#include "kmeans_centroid_sums.cuh" #include #include #include @@ -20,6 +21,7 @@ #include #include #include +#include #include #include #include @@ -42,7 +44,13 @@ #include #include #include +#include +#include #include +#include +#include +#include +#include #include #include @@ -121,6 +129,85 @@ struct WeightedDistanceOp { } }; +/** + * @brief A `cuda::mr::resource` that hands out temporary storage from one caller-owned buffer. + * + * The buffer only grows, and deallocation is a no-op, so repeated calls reuse the same memory + * instead of allocating (and possibly synchronizing) every time. Only one allocation may be live + * at a time. + */ +class workspace_buffer_resource { + public: + explicit workspace_buffer_resource(rmm::device_uvector& buffer) : buffer_{&buffer} {} + + void* allocate(cuda::stream_ref stream, size_t bytes, size_t alignment) + { + const size_t needed = bytes + alignment; + if (buffer_->size() < needed) { buffer_->resize(needed, stream); } + const auto ptr = reinterpret_cast(buffer_->data()); + return reinterpret_cast(raft::alignTo(ptr, alignment)); + } + void deallocate(cuda::stream_ref, void*, size_t, size_t) noexcept {} + void* allocate_sync(size_t bytes, size_t alignment) + { + return allocate(buffer_->stream(), bytes, alignment); + } + void deallocate_sync(void*, size_t, size_t) noexcept {} + bool operator==(const workspace_buffer_resource& other) const noexcept + { + return buffer_ == other.buffer_; + } + bool operator!=(const workspace_buffer_resource& other) const noexcept + { + return !(*this == other); + } + + private: + rmm::device_uvector* buffer_; +}; + +/** + * @brief Device-wide reduction that is bitwise reproducible from run to run. + * + * Requests `cuda::execution::determinism::run_to_run` explicitly (also cub's current default): + * cub combines the partial results in a fixed order on a given GPU, with no atomics. This is + * stronger than raft's mapThenSumReduce, which combines block results with atomics. + * + * `gpu_to_gpu` (an RFA-based sum that is identical across architectures) is not requested: the + * inputs of these reductions (distances from GEMM / cuTile TF32 kernels and the centroids derived + * from them) already differ between architectures, so a fit cannot be reproducible across GPUs. + * + * Temporary storage comes from `workspace`, which is resized as needed. + */ +template +void run_to_run_reduce(raft::resources const& handle, + InputIterator input, + IndexT size, + OutputT* output, + ReductionOpT reduction_op, + OutputT init, + rmm::device_uvector& workspace) +{ + auto env = + cuda::std::execution::env{cuda::stream_ref{raft::resource::get_cuda_stream(handle)}, + cuda::execution::require(cuda::execution::determinism::run_to_run), + cuda::std::execution::prop{cuda::mr::get_memory_resource, + workspace_buffer_resource{workspace}}}; + RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(input, output, size, reduction_op, init, env)); +} + +template +struct SquaredDiffOp { + const DataT* a; + const DataT* b; + + __device__ DataT operator()(IndexT offset) const + { + DataT diff = a[offset] - b[offset]; + return diff * diff; + } +}; + template void copyClusterLabels(raft::resources const& handle, const MinClusterAndDistanceResult& result, @@ -194,8 +281,14 @@ void weightSum( DataT wt_sum_h = DataT{0}; if constexpr (raft::is_device_mdspan_v) { - raft::linalg::mapThenSumReduce( - d_wt_sum.data_handle(), n_samples, raft::identity_op{}, stream.get(), weight.data_handle()); + rmm::device_uvector workspace(0, stream); + run_to_run_reduce(handle, + weight.data_handle(), + n_samples, + d_wt_sum.data_handle(), + cuda::std::plus<>{}, + DataT{0}, + workspace); if (check_positive) { raft::copy(&wt_sum_h, d_wt_sum.data_handle(), 1, stream); raft::resource::sync_stream(handle); @@ -270,29 +363,9 @@ void computeClusterCost(raft::resources const& handle, MainOpT main_op, ReductionOpT reduction_op) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); cuda::transform_iterator itr(input, main_op); - - size_t temp_storage_bytes = 0; - RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(nullptr, - temp_storage_bytes, - itr, - clusterCost.data_handle(), - size, - reduction_op, - OutputT(), - stream)); - - workspace.resize(temp_storage_bytes, stream); - - RAFT_CUDA_TRY(cub::DeviceReduce::Reduce(workspace.data(), - temp_storage_bytes, - itr, - clusterCost.data_handle(), - size, - reduction_op, - OutputT(), - stream)); + run_to_run_reduce( + handle, itr, size, clusterCost.data_handle(), reduction_op, OutputT(), workspace); } template & workspace, bool reset_sums = true) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); - auto n_samples = X.extent(0); - - workspace.resize(n_samples, stream); - - raft::linalg::reduce_rows_by_key(X.data_handle(), - X.extent(1), - cluster_labels, - sample_weights.data_handle(), - workspace.data(), - X.extent(0), - X.extent(1), - n_clusters, - centroid_sums.data_handle(), - stream, - reset_sums); - - raft::linalg::reduce_cols_by_key(sample_weights.data_handle(), - cluster_labels, - weight_per_cluster.data_handle(), - static_cast(1), - static_cast(n_samples), - n_clusters, - stream, - reset_sums); + sums_by_label::weighted_sums_by_label(handle, + X.data_handle(), + sample_weights.data_handle(), + cluster_labels, + X.extent(0), + X.extent(1), + n_clusters, + centroid_sums.data_handle(), + weight_per_cluster.data_handle(), + workspace, + reset_sums); } /** * @brief Finalize centroids by dividing accumulated sums by counts. @@ -707,21 +769,26 @@ void finalize_centroids(raft::resources const& handle, * @brief Compute the squared norm difference between two centroid sets. * * Writes sum((old_centroids - new_centroids)^2) into @p sqrd_norm_out. - * Used for convergence checking. Fully asynchronous — no stream sync. + * Used for convergence checking. Fully asynchronous — no stream sync. Temporary storage comes + * from @p workspace, which is resized as needed. */ template void compute_centroid_shift(raft::resources const& handle, raft::device_matrix_view old_centroids, raft::device_matrix_view new_centroids, - raft::device_scalar_view sqrd_norm_out) + raft::device_scalar_view sqrd_norm_out, + rmm::device_uvector& workspace) { - cudaStream_t stream = raft::resource::get_cuda_stream(handle).get(); - raft::linalg::mapThenSumReduce(sqrd_norm_out.data_handle(), - old_centroids.size(), - raft::sqdiff_op{}, - stream, - old_centroids.data_handle(), - new_centroids.data_handle()); + run_to_run_reduce( + handle, + cuda::transform_iterator( + cuda::counting_iterator(0), + SquaredDiffOp{old_centroids.data_handle(), new_centroids.data_handle()}), + static_cast(old_centroids.size()), + sqrd_norm_out.data_handle(), + cuda::std::plus<>{}, + DataT{0}, + workspace); } /** diff --git a/cpp/src/cluster/detail/kmeans_mg.cuh b/cpp/src/cluster/detail/kmeans_mg.cuh index a86c30a087..4b25221f63 100644 --- a/cpp/src/cluster/detail/kmeans_mg.cuh +++ b/cpp/src/cluster/detail/kmeans_mg.cuh @@ -605,7 +605,8 @@ void mnmg_fit( dev_res, rank_centroids_const, raft::make_const_mdspan(new_centroids.view()), - sqrd_norm_error_dev.view()); + sqrd_norm_error_dev.view(), + workspace); raft::copy( rank_centroids.data_handle(), new_centroids.data_handle(), n_clusters * n_features, stream); diff --git a/cpp/tests/cluster/kmeans.cu b/cpp/tests/cluster/kmeans.cu index 23fc248917..a6e5b4c0f6 100644 --- a/cpp/tests/cluster/kmeans.cu +++ b/cpp/tests/cluster/kmeans.cu @@ -23,6 +23,7 @@ #include +#include #include #include @@ -495,6 +496,82 @@ class KmeansFitBatchedTest : public ::testing::TestWithParam(handle, n_clusters, n_features); + raft::random::RngState rng(p.rng_state.seed); + raft::random::uniform(handle, rng, d_init.data_handle(), n_clusters * n_features, T(-1), T(1)); + + auto d_sw = d_sw_view(); + std::optional> h_sw = std::nullopt; + auto h_sample_weight = raft::make_host_vector(testparams.weighted ? n_samples : 0); + if (testparams.weighted) { + std::fill_n(h_sample_weight.data_handle(), n_samples, T(1)); + h_sw = std::make_optional(raft::make_const_mdspan(h_sample_weight.view())); + } + + constexpr int n_runs = 3; + for (int host : {0, 1}) { + std::vector> centroids(n_runs, std::vector(n_clusters * n_features)); + std::vector inertias(n_runs); + std::vector n_iters(n_runs); + for (int r = 0; r < n_runs; ++r) { + auto d_c = raft::make_device_matrix(handle, n_clusters, n_features); + raft::copy(d_c.data_handle(), d_init.data_handle(), n_clusters * n_features, stream); + T inertia = 0; + if (host) { + auto pb = p; + pb.device_buffer_samples = testparams.device_buffer_samples; + int64_t n_iter = 0; + cuvs::cluster::kmeans::fit( + handle, + pb, + raft::make_const_mdspan(h_X->view()), + h_sw, + raft::make_device_matrix_view(d_c.data_handle(), n_clusters, n_features), + raft::make_host_scalar_view(&inertia), + raft::make_host_scalar_view(&n_iter)); + n_iters[r] = n_iter; + } else { + int n_iter = 0; + cuvs::cluster::kmeans::fit(handle, + p, + raft::make_const_mdspan(d_X->view()), + d_sw, + d_c.view(), + raft::make_host_scalar_view(&inertia), + raft::make_host_scalar_view(&n_iter)); + n_iters[r] = n_iter; + } + raft::copy(centroids[r].data(), d_c.data_handle(), n_clusters * n_features, stream); + raft::resource::sync_stream(handle, stream); + inertias[r] = inertia; + } + for (int r = 1; r < n_runs; ++r) { + EXPECT_EQ(n_iters[0], n_iters[r]) << (host ? "host" : "device") << " run " << r; + EXPECT_EQ(inertias[0], inertias[r]) << (host ? "host" : "device") << " run " << r; + EXPECT_EQ( + 0, std::memcmp(centroids[0].data(), centroids[r].data(), centroids[0].size() * sizeof(T))) + << (host ? "host" : "device") << " run " << r << ": centroids differ"; + } + } + } + T fitHostWithInitSize(int64_t init_size_value) { int n_features = testparams.n_col; @@ -709,6 +786,18 @@ TEST_P(KmeansFitBatchedTestD, Result) runZeroCost(); } +TEST_P(KmeansFitBatchedTestF, Reproducible) +{ + prepareBlobInputs(); + runReproducibleCheck(); +} + +TEST_P(KmeansFitBatchedTestD, Reproducible) +{ + prepareBlobInputs(); + runReproducibleCheck(); +} + INSTANTIATE_TEST_CASE_P(KmeansFitBatchedTests, KmeansFitBatchedTestF, ::testing::ValuesIn(batched_inputsf2));