Add rust bindings for kmeans

rust/cuvs-sys/cuvs_c_wrapper.h

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2024, NVIDIA CORPORATION.
+ * Copyright (c) 2024-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -17,6 +17,7 @@
 // wrapper file containing all the C-API's we should automatically be creating rust
 // bindings for
 #include <cuvs/core/c_api.h>
+#include <cuvs/cluster/kmeans.h>
 #include <cuvs/distance/pairwise_distance.h>
 #include <cuvs/neighbors/brute_force.h>
 #include <cuvs/neighbors/ivf_flat.h>

rust/cuvs/src/cluster/kmeans/mod.rs

@@ -0,0 +1,212 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+//! Kmeans clustering API's
+//!
+//! Example:
+//! ```
+//!
+//! use cuvs::cluster::kmeans;
+//! use cuvs::{ManagedTensor, Resources, Result};
+//!
+//! use ndarray_rand::rand_distr::Uniform;
+//! use ndarray_rand::RandomExt;
+//!
+//! fn kmeans_example() -> Result<()> {
+//!     let res = Resources::new()?;
+//!
+//!     // Create a new random dataset to index
+//!     let n_datapoints = 65536;
+//!     let n_features = 512;
+//!     let n_clusters = 8;
+//!     let dataset =
+//!         ndarray::Array::<f32, _>::random((n_datapoints, n_features), Uniform::new(0., 1.0));
+//!     let dataset = ManagedTensor::from(&dataset).to_device(&res)?;
+//!
+//!     let centroids_host = ndarray::Array::<f32, _>::zeros((n_clusters, n_features));
+//!     let mut centroids = ManagedTensor::from(&centroids_host).to_device(&res)?;
+//!
+//!     // find the centroids with the kmeans index
+//!     let kmeans_params = kmeans::Params::new()?.set_n_clusters(n_clusters as i32);
+//!     let (inertia, n_iter) = kmeans::fit(&res, &kmeans_params, &dataset, &None, &mut centroids)?;
+//!     Ok(())
+//! }
+//! ```
+
+mod params;
+
+pub use params::Params;
+
+use crate::dlpack::ManagedTensor;
+use crate::error::{check_cuvs, Result};
+use crate::resources::Resources;
+
+/// Find clusters with the k-means algorithm
+///
+/// # Arguments
+///
+/// * `res` - Resources to use
+/// * `params` - Parameters to use to fit KMeans model
+/// * `x` - A matrix in device memory - shape (m, k)
+/// * `sample_weight` - Optional device matrix shape (n_clusters, 1)
+/// * `centroids` - Output device matrix, that has the centroids for each cluster
+///   shape (n_clusters, k)
+pub fn fit(
+    res: &Resources,
+    params: &Params,
+    x: &ManagedTensor,
+    sample_weight: &Option<ManagedTensor>,
+    centroids: &mut ManagedTensor,
+) -> Result<(f64, i32)> {
+    let mut inertia: f64 = 0.0;
+    let mut niter: i32 = 0;
+
+    unsafe {
+        let sample_weight_dlpack = match sample_weight {
+            Some(tensor) => tensor.as_ptr(),
+            None => std::ptr::null_mut(),
+        };
+        check_cuvs(ffi::cuvsKMeansFit(
+            res.0,
+            params.0,
+            x.as_ptr(),
+            sample_weight_dlpack,
+            centroids.as_ptr(),
+            &mut inertia as *mut f64,
+            &mut niter as *mut i32,
+        ))?;
+    }
+    Ok((inertia, niter))
+}
+
+/// Predict clusters with the k-means algorithm
+///
+/// # Arguments
+///
+/// * `res` - Resources to use
+/// * `params` - Parameters to use to fit KMeans model
+/// * `x` - Input matrix in device memory - shape (m, k)
+/// * `sample_weight` - Optional device matrix shape (n_clusters, 1)
+/// * `centroids` - Centroids calculated by fit in device memory, shape (n_clusters, k)
+/// * `labels` - preallocated CUDA array interface matrix shape (m, 1) to hold the output labels
+/// * `normalize_weight` - whether or not to normalize the weights
+pub fn predict(
+    res: &Resources,
+    params: &Params,
+    x: &ManagedTensor,
+    sample_weight: &Option<ManagedTensor>,
+    centroids: &ManagedTensor,
+    labels: &mut ManagedTensor,
+    normalize_weight: bool,
+) -> Result<f64> {
+    let mut inertia: f64 = 0.0;
+
+    unsafe {
+        let sample_weight_dlpack = match sample_weight {
+            Some(tensor) => tensor.as_ptr(),
+            None => std::ptr::null_mut(),
+        };
+        check_cuvs(ffi::cuvsKMeansPredict(
+            res.0,
+            params.0,
+            x.as_ptr(),
+            sample_weight_dlpack,
+            centroids.as_ptr(),
+            labels.as_ptr(),
+            normalize_weight,
+            &mut inertia as *mut f64,
+        ))?;
+    }
+    Ok(inertia)
+}
+
+/// Compute cluster cost given an input matrix and existing centroids
+/// # Arguments
+///
+/// * `res` - Resources to use
+/// * `x` - Input matrix in device memory - shape (m, k)
+/// * `centroids` - Centroids calculated by fit in device memory, shape (n_clusters, k)
+pub fn cluster_cost(res: &Resources, x: &ManagedTensor, centroids: &ManagedTensor) -> Result<f64> {
+    let mut inertia: f64 = 0.0;
+
+    unsafe {
+        check_cuvs(ffi::cuvsKMeansClusterCost(
+            res.0,
+            x.as_ptr(),
+            centroids.as_ptr(),
+            &mut inertia as *mut f64,
+        ))?;
+    }
+    Ok(inertia)
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use ndarray_rand::rand_distr::Uniform;
+    use ndarray_rand::RandomExt;
+
+    #[test]
+    fn test_kmeans() {
+        let res = Resources::new().unwrap();
+
+        let n_clusters = 4;
+
+        // Create a new random dataset to index
+        let n_datapoints = 256;
+        let n_features = 16;
+        let dataset =
+            ndarray::Array::<f32, _>::random((n_datapoints, n_features), Uniform::new(0., 1.0));
+        let dataset = ManagedTensor::from(&dataset).to_device(&res).unwrap();
+
+        let centroids_host = ndarray::Array::<f32, _>::zeros((n_clusters, n_features));
+        let mut centroids = ManagedTensor::from(&centroids_host)
+            .to_device(&res)
+            .unwrap();
+
+        let params = Params::new().unwrap().set_n_clusters(n_clusters as i32);
+
+        // compute the inertia, before fitting centroids
+        let original_inertia = cluster_cost(&res, &dataset, &centroids).unwrap();
+
+        // fit the centroids, make sure that inertia has gone down
+        let (inertia, n_iter) = fit(&res, &params, &dataset, &None, &mut centroids).unwrap();
+
+        assert!(inertia < original_inertia);
+        assert!(n_iter >= 1);
+
+        let mut labels_host = ndarray::Array::<i32, _>::zeros((n_clusters,));
+        let mut labels = ManagedTensor::from(&labels_host).to_device(&res).unwrap();
+
+        // make sure the prediction for each centroid is the centroid itself
+        predict(
+            &res,
+            &params,
+            &centroids,
+            &None,
+            &centroids,
+            &mut labels,
+            false,
+        )
+        .unwrap();
+
+        labels.to_host(&res, &mut labels_host).unwrap();
+        assert_eq!(labels_host[[0,]], 0);
+        assert_eq!(labels_host[[1,]], 1);
+        assert_eq!(labels_host[[2,]], 2);
+        assert_eq!(labels_host[[3,]], 3);
+    }
+}

rust/cuvs/src/cluster/kmeans/params.rs

@@ -0,0 +1,152 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+use crate::distance_type::DistanceType;
+use crate::error::{check_cuvs, Result};
+use std::fmt;
+use std::io::{stderr, Write};
+
+pub struct Params(pub ffi::cuvsKMeansParams_t);
+
+impl Params {
+    /// Returns a new Params
+    pub fn new() -> Result<Params> {
+        unsafe {
+            let mut params = std::mem::MaybeUninit::<ffi::cuvsKMeansParams_t>::uninit();
+            check_cuvs(ffi::cuvsKMeansParamsCreate(params.as_mut_ptr()))?;
+            Ok(Params(params.assume_init()))
+        }
+    }
+
+    /// DistanceType to use for fitting kmeans
+    pub fn set_metric(self, metric: DistanceType) -> Params {
+        unsafe {
+            (*self.0).metric = metric;
+        }
+        self
+    }
+
+    /// The number of clusters to form as well as the number of centroids to generate (default:8).
+    pub fn set_n_clusters(self, n_clusters: i32) -> Params {
+        unsafe {
+            (*self.0).n_clusters = n_clusters;
+        }
+        self
+    }
+
+    /// Maximum number of iterations of the k-means algorithm for a single run.
+    pub fn set_max_iter(self, max_iter: i32) -> Params {
+        unsafe {
+            (*self.0).max_iter = max_iter;
+        }
+        self
+    }
+
+    /// Relative tolerance with regards to inertia to declare convergence.
+    pub fn set_tol(self, tol: f64) -> Params {
+        unsafe {
+            (*self.0).tol = tol;
+        }
+        self
+    }
+
+    /// Number of instance k-means algorithm will be run with different seeds.
+    pub fn set_n_init(self, n_init: i32) -> Params {
+        unsafe {
+            (*self.0).n_init = n_init;
+        }
+        self
+    }
+
+    /// Oversampling factor for use in the k-means|| algorithm
+    pub fn set_oversampling_factor(self, oversampling_factor: f64) -> Params {
+        unsafe {
+            (*self.0).oversampling_factor = oversampling_factor;
+        }
+        self
+    }
+
+    /**
+     * batch_samples and batch_centroids are used to tile 1NN computation which is
+     * useful to optimize/control the memory footprint
+     * Default tile is [batch_samples x n_clusters] i.e. when batch_centroids is 0
+     * then don't tile the centroids.
+     */
+    pub fn set_batch_samples(self, batch_samples: i32) -> Params {
+        unsafe {
+            (*self.0).batch_samples = batch_samples;
+        }
+        self
+    }
+    /// if 0 then batch_centroids = n_clusters
+    pub fn set_batch_centroids(self, batch_centroids: i32) -> Params {
+        unsafe {
+            (*self.0).batch_centroids = batch_centroids;
+        }
+        self
+    }
+
+    /// Whether to use hierarchical (balanced) kmeans or not
+    pub fn set_hierarchical(self, hierarchical: bool) -> Params {
+        unsafe {
+            (*self.0).hierarchical = hierarchical;
+        }
+        self
+    }
+
+    /// For hierarchical k-means , defines the number of training iterations
+    pub fn set_hierarchical_n_iters(self, hierarchical_n_iters: i32) -> Params {
+        unsafe {
+            (*self.0).hierarchical_n_iters = hierarchical_n_iters;
+        }
+        self
+    }
+}
+
+impl fmt::Debug for Params {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        // custom debug trait here, default value will show the pointer address
+        // for the inner params object which isn't that useful.
+        write!(f, "Params({:?})", unsafe { *self.0 })
+    }
+}
+
+impl Drop for Params {
+    fn drop(&mut self) {
+        if let Err(e) = check_cuvs(unsafe { ffi::cuvsKMeansParamsDestroy(self.0) }) {
+            write!(stderr(), "failed to call cuvsKMeansParamsDestroy {:?}", e)
+                .expect("failed to write to stderr");
+        }
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_params() {
+        let params = Params::new()
+            .unwrap()
+            .set_n_clusters(128)
+            .set_hierarchical(true);
+
+        unsafe {
+            assert_eq!((*params.0).n_clusters, 128);
+            assert_eq!((*params.0).hierarchical, true);
+        }
+    }
+}

rust/cuvs/src/cluster/mod.rs

@@ -0,0 +1,17 @@
+/*
+ * Copyright (c) 2025, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+pub mod kmeans;