|
| 1 | +/* |
| 2 | + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. |
| 3 | + * SPDX-License-Identifier: Apache-2.0 |
| 4 | + */ |
| 5 | + |
| 6 | +#include <cstdint> |
| 7 | +#include <optional> |
| 8 | + |
| 9 | +#include <dlpack/dlpack.h> |
| 10 | + |
| 11 | +#include <cuvs/cluster/kmeans.hpp> |
| 12 | +#include <cuvs/cluster/mg_kmeans.h> |
| 13 | +#include <cuvs/core/c_api.h> |
| 14 | + |
| 15 | +#include <raft/core/device_mdarray.hpp> |
| 16 | +#include <raft/core/resource/cuda_stream.hpp> |
| 17 | +#include <raft/core/resource/multi_gpu.hpp> |
| 18 | +#include <raft/core/resources.hpp> |
| 19 | +#include <raft/util/cudart_utils.hpp> |
| 20 | + |
| 21 | +#include "../core/exceptions.hpp" |
| 22 | +#include "../core/interop.hpp" |
| 23 | + |
| 24 | +namespace { |
| 25 | + |
| 26 | +template <typename ParamsT> |
| 27 | +cuvs::cluster::kmeans::params convert_params(const ParamsT& params) |
| 28 | +{ |
| 29 | + auto kmeans_params = cuvs::cluster::kmeans::params(); |
| 30 | + kmeans_params.metric = static_cast<cuvs::distance::DistanceType>(params.metric); |
| 31 | + kmeans_params.init = static_cast<cuvs::cluster::kmeans::params::InitMethod>(params.init); |
| 32 | + kmeans_params.n_clusters = params.n_clusters; |
| 33 | + kmeans_params.max_iter = params.max_iter; |
| 34 | + kmeans_params.tol = params.tol; |
| 35 | + kmeans_params.n_init = params.n_init; |
| 36 | + kmeans_params.oversampling_factor = params.oversampling_factor; |
| 37 | + kmeans_params.batch_samples = params.batch_samples; |
| 38 | + kmeans_params.batch_centroids = params.batch_centroids; |
| 39 | + kmeans_params.init_size = params.init_size; |
| 40 | + kmeans_params.streaming_batch_size = params.streaming_batch_size; |
| 41 | + return kmeans_params; |
| 42 | +} |
| 43 | + |
| 44 | +void validate_host_tensor(DLManagedTensor* tensor, const char* name) |
| 45 | +{ |
| 46 | + RAFT_EXPECTS(tensor != nullptr, "%s must not be NULL", name); |
| 47 | + auto dl_tensor = tensor->dl_tensor; |
| 48 | + RAFT_EXPECTS(dl_tensor.data != nullptr, "%s data must not be NULL", name); |
| 49 | + RAFT_EXPECTS(dl_tensor.shape != nullptr, "%s shape must not be NULL", name); |
| 50 | + RAFT_EXPECTS( |
| 51 | + cuvs::core::is_dlpack_host_compatible(dl_tensor), "%s must be host accessible", name); |
| 52 | + RAFT_EXPECTS(dl_tensor.device.device_type != kDLCUDA, "%s must reside in host memory", name); |
| 53 | + RAFT_EXPECTS(cuvs::core::is_c_contiguous(tensor), "%s must be C-contiguous", name); |
| 54 | +} |
| 55 | + |
| 56 | +bool dtype_equal(const DLTensor& lhs, const DLTensor& rhs) |
| 57 | +{ |
| 58 | + return lhs.dtype.code == rhs.dtype.code && lhs.dtype.bits == rhs.dtype.bits && |
| 59 | + lhs.dtype.lanes == rhs.dtype.lanes; |
| 60 | +} |
| 61 | + |
| 62 | +void validate_float_dtype(const DLTensor& tensor, const char* name) |
| 63 | +{ |
| 64 | + RAFT_EXPECTS( |
| 65 | + tensor.dtype.code == kDLFloat && (tensor.dtype.bits == 32 || tensor.dtype.bits == 64), |
| 66 | + "%s must have dtype float32 or float64", |
| 67 | + name); |
| 68 | + RAFT_EXPECTS(tensor.dtype.lanes == 1, "%s must have one DLPack lane", name); |
| 69 | +} |
| 70 | + |
| 71 | +template <typename ParamsT> |
| 72 | +void validate_inputs(const ParamsT& params, |
| 73 | + DLManagedTensor* X_tensor, |
| 74 | + DLManagedTensor* sample_weight_tensor, |
| 75 | + DLManagedTensor* centroids_tensor) |
| 76 | +{ |
| 77 | + RAFT_EXPECTS(params.n_clusters > 0, "n_clusters must be positive"); |
| 78 | + RAFT_EXPECTS(!params.hierarchical, "hierarchical kmeans is not supported by SNMG kmeans"); |
| 79 | + |
| 80 | + validate_host_tensor(X_tensor, "X"); |
| 81 | + validate_host_tensor(centroids_tensor, "centroids"); |
| 82 | + |
| 83 | + auto X = X_tensor->dl_tensor; |
| 84 | + auto centroids = centroids_tensor->dl_tensor; |
| 85 | + |
| 86 | + RAFT_EXPECTS(X.ndim == 2, "X must be a 2D tensor"); |
| 87 | + RAFT_EXPECTS(centroids.ndim == 2, "centroids must be a 2D tensor"); |
| 88 | + RAFT_EXPECTS(X.shape[0] > 0, "X must have at least one row"); |
| 89 | + RAFT_EXPECTS(X.shape[1] > 0, "X must have at least one column"); |
| 90 | + RAFT_EXPECTS(centroids.shape[0] == params.n_clusters, |
| 91 | + "centroids row count must equal n_clusters"); |
| 92 | + RAFT_EXPECTS(centroids.shape[1] == X.shape[1], |
| 93 | + "centroids column count must equal X column count"); |
| 94 | + |
| 95 | + validate_float_dtype(X, "X"); |
| 96 | + RAFT_EXPECTS(dtype_equal(X, centroids), "centroids dtype must match X dtype"); |
| 97 | + |
| 98 | + if (sample_weight_tensor != nullptr) { |
| 99 | + validate_host_tensor(sample_weight_tensor, "sample_weight"); |
| 100 | + auto sample_weight = sample_weight_tensor->dl_tensor; |
| 101 | + RAFT_EXPECTS(sample_weight.ndim == 1, "sample_weight must be a 1D tensor"); |
| 102 | + RAFT_EXPECTS(sample_weight.shape[0] == X.shape[0], |
| 103 | + "sample_weight length must equal X row count"); |
| 104 | + RAFT_EXPECTS(dtype_equal(X, sample_weight), "sample_weight dtype must match X dtype"); |
| 105 | + } |
| 106 | +} |
| 107 | + |
| 108 | +template <typename T, typename ParamsT, typename IdxT = int64_t> |
| 109 | +void fit_snmg(cuvsResources_t res, |
| 110 | + const ParamsT& params, |
| 111 | + DLManagedTensor* X_tensor, |
| 112 | + DLManagedTensor* sample_weight_tensor, |
| 113 | + DLManagedTensor* centroids_tensor, |
| 114 | + double* inertia, |
| 115 | + int* n_iter) |
| 116 | +{ |
| 117 | + auto res_ptr = reinterpret_cast<raft::resources*>(res); |
| 118 | + RAFT_EXPECTS(res_ptr != nullptr, "res must not be NULL"); |
| 119 | + RAFT_EXPECTS(raft::resource::is_multi_gpu(*res_ptr), |
| 120 | + "cuvsMultiGpuKMeansFit requires a MultiGpuResources handle"); |
| 121 | + |
| 122 | + auto X = X_tensor->dl_tensor; |
| 123 | + auto centroids = centroids_tensor->dl_tensor; |
| 124 | + |
| 125 | + auto n_samples = static_cast<IdxT>(X.shape[0]); |
| 126 | + auto n_features = static_cast<IdxT>(X.shape[1]); |
| 127 | + auto n_clusters = static_cast<IdxT>(params.n_clusters); |
| 128 | + |
| 129 | + auto X_view = raft::make_host_matrix_view<T const, IdxT>( |
| 130 | + reinterpret_cast<T const*>(X.data), n_samples, n_features); |
| 131 | + |
| 132 | + std::optional<raft::host_vector_view<T const, IdxT>> sample_weight; |
| 133 | + if (sample_weight_tensor != nullptr) { |
| 134 | + auto sw = sample_weight_tensor->dl_tensor; |
| 135 | + sample_weight = |
| 136 | + raft::make_host_vector_view<T const, IdxT>(reinterpret_cast<T const*>(sw.data), n_samples); |
| 137 | + } |
| 138 | + |
| 139 | + auto const& rank0_res = raft::resource::set_current_device_to_rank(*res_ptr, 0); |
| 140 | + auto stream = raft::resource::get_cuda_stream(rank0_res); |
| 141 | + auto d_centroids = raft::make_device_matrix<T, IdxT>(rank0_res, n_clusters, n_features); |
| 142 | + auto n_centroid_values = n_clusters * n_features; |
| 143 | + |
| 144 | + if (params.init == Array) { |
| 145 | + raft::update_device(d_centroids.data_handle(), |
| 146 | + reinterpret_cast<T const*>(centroids.data), |
| 147 | + n_centroid_values, |
| 148 | + stream); |
| 149 | + raft::resource::sync_stream(rank0_res, stream); |
| 150 | + } |
| 151 | + |
| 152 | + T inertia_temp = T{0}; |
| 153 | + IdxT n_iter_temp = IdxT{0}; |
| 154 | + auto kmeans_params = convert_params(params); |
| 155 | + cuvs::cluster::kmeans::fit(*res_ptr, |
| 156 | + kmeans_params, |
| 157 | + X_view, |
| 158 | + sample_weight, |
| 159 | + d_centroids.view(), |
| 160 | + raft::make_host_scalar_view<T>(&inertia_temp), |
| 161 | + raft::make_host_scalar_view<IdxT>(&n_iter_temp)); |
| 162 | + |
| 163 | + raft::update_host( |
| 164 | + reinterpret_cast<T*>(centroids.data), d_centroids.data_handle(), n_centroid_values, stream); |
| 165 | + raft::resource::sync_stream(rank0_res, stream); |
| 166 | + |
| 167 | + *inertia = static_cast<double>(inertia_temp); |
| 168 | + *n_iter = static_cast<int>(n_iter_temp); |
| 169 | +} |
| 170 | + |
| 171 | +template <typename ParamsT> |
| 172 | +void dispatch_fit(cuvsResources_t res, |
| 173 | + ParamsT params, |
| 174 | + DLManagedTensor* X, |
| 175 | + DLManagedTensor* sample_weight, |
| 176 | + DLManagedTensor* centroids, |
| 177 | + double* inertia, |
| 178 | + int* n_iter) |
| 179 | +{ |
| 180 | + RAFT_EXPECTS(res != 0, "res must not be NULL"); |
| 181 | + RAFT_EXPECTS(params != nullptr, "params must not be NULL"); |
| 182 | + RAFT_EXPECTS(inertia != nullptr, "inertia must not be NULL"); |
| 183 | + RAFT_EXPECTS(n_iter != nullptr, "n_iter must not be NULL"); |
| 184 | + |
| 185 | + validate_inputs(*params, X, sample_weight, centroids); |
| 186 | + |
| 187 | + auto dataset = X->dl_tensor; |
| 188 | + if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) { |
| 189 | + fit_snmg<float>(res, *params, X, sample_weight, centroids, inertia, n_iter); |
| 190 | + } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 64) { |
| 191 | + fit_snmg<double>(res, *params, X, sample_weight, centroids, inertia, n_iter); |
| 192 | + } else { |
| 193 | + RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d", |
| 194 | + dataset.dtype.code, |
| 195 | + dataset.dtype.bits); |
| 196 | + } |
| 197 | +} |
| 198 | + |
| 199 | +} // namespace |
| 200 | + |
| 201 | +extern "C" cuvsError_t cuvsMultiGpuKMeansFit(cuvsResources_t res, |
| 202 | + cuvsKMeansParams_v2_t params, |
| 203 | + DLManagedTensor* X, |
| 204 | + DLManagedTensor* sample_weight, |
| 205 | + DLManagedTensor* centroids, |
| 206 | + double* inertia, |
| 207 | + int* n_iter) |
| 208 | +{ |
| 209 | + return cuvs::core::translate_exceptions( |
| 210 | + [=] { dispatch_fit(res, params, X, sample_weight, centroids, inertia, n_iter); }); |
| 211 | +} |
0 commit comments