Introduce fused GEMM + 1-NN primitive using cuTile

c/src/cluster/kmeans.cpp

@@ -4,6 +4,7 @@
  */
 
 #include <cstdint>
+#include <limits>
 
 #include <dlpack/dlpack.h>
 
@@ -41,6 +42,56 @@ cuvs::cluster::kmeans::balanced_params convert_balanced_params(const cuvsKMeansP
   return kmeans_params;
 }
 
+constexpr int64_t kKMeansInt32IndexMax = std::numeric_limits<int32_t>::max();
+
+bool dlpack_shape_exceeds_int32_index(const DLTensor& tensor)
+{
+  for (int i = 0; i < tensor.ndim; ++i) {
+    if (tensor.shape[i] > kKMeansInt32IndexMax) { return true; }
+  }
+  return false;
+}
+
+bool kmeans_tensor_shapes_use_int64_index(DLManagedTensor* X, DLManagedTensor* centroids)
+{
+  if (dlpack_shape_exceeds_int32_index(X->dl_tensor)) { return true; }
+  if (dlpack_shape_exceeds_int32_index(centroids->dl_tensor)) { return true; }
+  return false;
+}
+
+bool kmeans_fit_uses_int64_index(DLManagedTensor* X,
+                                 DLManagedTensor* centroids,
+                                 int n_clusters)
+{
+  if (static_cast<int64_t>(n_clusters) > kKMeansInt32IndexMax) { return true; }
+  return kmeans_tensor_shapes_use_int64_index(X, centroids);
+}
+
+bool kmeans_labels_use_int64_index(const DLTensor& labels)
+{
+  return labels.dtype.code == kDLInt && labels.dtype.bits == 64;
+}
+
+void validate_kmeans_labels_dtype(const DLTensor& labels)
+{
+  if (labels.dtype.code == kDLInt && (labels.dtype.bits == 32 || labels.dtype.bits == 64)) {
+    return;
+  }
+  RAFT_FAIL("Unsupported labels DLtensor dtype: %d and bits: %d",
+            labels.dtype.code,
+            labels.dtype.bits);
+}
+
+bool kmeans_predict_uses_int64_index(DLManagedTensor* X,
+                                     DLManagedTensor* centroids,
+                                     DLManagedTensor* labels,
+                                     int n_clusters)
+{
+  validate_kmeans_labels_dtype(labels->dl_tensor);
+  if (kmeans_labels_use_int64_index(labels->dl_tensor)) { return true; }
+  return kmeans_fit_uses_int64_index(X, centroids, n_clusters);
+}
+
 template <typename T, typename IdxT = int64_t>
 void _fit(cuvsResources_t res,
           const cuvsKMeansParams& params,
@@ -54,8 +105,10 @@ void _fit(cuvsResources_t res,
   auto res_ptr = reinterpret_cast<raft::resources*>(res);
 
   if (!cuvs::core::is_dlpack_device_compatible(X)) {
-    auto n_samples  = static_cast<IdxT>(X.shape[0]);
-    auto n_features = static_cast<IdxT>(X.shape[1]);
+    // Host fit overloads are only exposed with int64_t index types.
+    using HostIdxT = int64_t;
+    auto n_samples  = static_cast<HostIdxT>(X.shape[0]);
+    auto n_features = static_cast<HostIdxT>(X.shape[1]);
 
     if (params.hierarchical) {
       RAFT_FAIL("hierarchical kmeans is not supported with host data");
@@ -66,24 +119,24 @@ void _fit(cuvsResources_t res,
       RAFT_FAIL("centroids must be on device memory");
     }
 
-    auto X_view = raft::make_host_matrix_view<T const, IdxT>(
+    auto X_view = raft::make_host_matrix_view<T const, HostIdxT>(
       reinterpret_cast<T const*>(X.data), n_samples, n_features);
     auto centroids_view =
-      cuvs::core::from_dlpack<raft::device_matrix_view<T, IdxT, raft::row_major>>(
+      cuvs::core::from_dlpack<raft::device_matrix_view<T, HostIdxT, raft::row_major>>(
         centroids_tensor);
 
-    std::optional<raft::host_vector_view<T const, IdxT>> sample_weight;
+    std::optional<raft::host_vector_view<T const, HostIdxT>> sample_weight;
     if (sample_weight_tensor != NULL) {
       auto sw = sample_weight_tensor->dl_tensor;
       if (!cuvs::core::is_dlpack_host_compatible(sw)) {
         RAFT_FAIL("sample_weight must be host accessible when X is on host");
       }
-      sample_weight = raft::make_host_vector_view<T const, IdxT>(
+      sample_weight = raft::make_host_vector_view<T const, HostIdxT>(
         reinterpret_cast<T const*>(sw.data), n_samples);
     }
 
     T inertia_temp;
-    IdxT n_iter_temp;
+    HostIdxT n_iter_temp;
 
     auto kmeans_params = convert_params(params);
     cuvs::cluster::kmeans::fit(*res_ptr,
@@ -92,7 +145,7 @@ void _fit(cuvsResources_t res,
                                sample_weight,
                                centroids_view,
                                raft::make_host_scalar_view<T>(&inertia_temp),
-                               raft::make_host_scalar_view<IdxT>(&n_iter_temp));
+                               raft::make_host_scalar_view<HostIdxT>(&n_iter_temp));
 
     *inertia = inertia_temp;
     *n_iter  = n_iter_temp;
@@ -109,10 +162,20 @@ void _fit(cuvsResources_t res,
       if constexpr (std::is_same_v<T, double>) {
         RAFT_FAIL("float64 is an unsupported dtype for hierarchical kmeans");
       } else {
-        auto kmeans_params = convert_balanced_params(params);
+        // Balanced fit overloads are only exposed with int64_t index types.
+        using BalancedIdxT = int64_t;
+        using balanced_const_mdspan_type =
+          raft::device_matrix_view<T const, BalancedIdxT, raft::row_major>;
+        using balanced_mdspan_type = raft::device_matrix_view<T, BalancedIdxT, raft::row_major>;
+        auto kmeans_params         = convert_balanced_params(params);
         T inertia_temp;
         auto inertia_view = raft::make_host_scalar_view<T>(&inertia_temp);
-        cuvs::cluster::kmeans::fit(*res_ptr, kmeans_params, cuvs::core::from_dlpack<const_mdspan_type>(X_tensor), cuvs::core::from_dlpack<mdspan_type>(centroids_tensor), std::make_optional(inertia_view));
+        cuvs::cluster::kmeans::fit(
+          *res_ptr,
+          kmeans_params,
+          cuvs::core::from_dlpack<balanced_const_mdspan_type>(X_tensor),
+          cuvs::core::from_dlpack<balanced_mdspan_type>(centroids_tensor),
+          std::make_optional(inertia_view));
         *inertia = inertia_temp;
         *n_iter  = params.hierarchical_n_iters;
       }
@@ -164,13 +227,22 @@ void _predict(cuvsResources_t res,
 
       if constexpr (std::is_same_v<T, double>) {
         RAFT_FAIL("float64 is an unsupported dtype for hierarchical kmeans");
+      } else if constexpr (!std::is_same_v<LabelsT, int32_t>) {
+        RAFT_FAIL("int64 labels are unsupported for hierarchical kmeans");
       } else {
+        // Balanced predict overloads are only exposed with int64_t index types and int32 labels.
+        using BalancedIdxT = int64_t;
+        using balanced_const_mdspan_type =
+          raft::device_matrix_view<T const, BalancedIdxT, raft::row_major>;
+        using balanced_labels_mdspan_type =
+          raft::device_vector_view<int32_t, BalancedIdxT, raft::row_major>;
         auto kmeans_params = convert_balanced_params(params);
-        cuvs::cluster::kmeans::predict(*res_ptr,
-                                       kmeans_params,
-                                       cuvs::core::from_dlpack<const_mdspan_type>(X_tensor),
-                                       cuvs::core::from_dlpack<const_mdspan_type>(centroids_tensor),
-                                       cuvs::core::from_dlpack<labels_mdspan_type>(labels_tensor));
+        cuvs::cluster::kmeans::predict(
+          *res_ptr,
+          kmeans_params,
+          cuvs::core::from_dlpack<balanced_const_mdspan_type>(X_tensor),
+          cuvs::core::from_dlpack<balanced_const_mdspan_type>(centroids_tensor),
+          cuvs::core::from_dlpack<balanced_labels_mdspan_type>(labels_tensor));
         *inertia = 0;
       }
     } else {
@@ -196,7 +268,7 @@ void _predict(cuvsResources_t res,
   }
 }
 
-template <typename T, typename IdxT = int32_t>
+template <typename T, typename IdxT>
 void _cluster_cost(cuvsResources_t res,
                    DLManagedTensor* X_tensor,
                    DLManagedTensor* centroids_tensor,
@@ -259,10 +331,24 @@ extern "C" cuvsError_t cuvsKMeansFit(cuvsResources_t res,
 {
   return cuvs::core::translate_exceptions([=] {
     auto dataset = X->dl_tensor;
+    const bool use_int64_index =
+      kmeans_fit_uses_int64_index(X, centroids, params->n_clusters);
     if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) {
-      _fit<float>(res, *params, X, sample_weight, centroids, inertia, n_iter);
+      if (use_int64_index) {
+        _fit<float, int64_t>(
+          res, *params, X, sample_weight, centroids, inertia, n_iter);
+      } else {
+        _fit<float, int32_t>(
+          res, *params, X, sample_weight, centroids, inertia, n_iter);
+      }
     } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 64) {
-      _fit<double>(res, *params, X, sample_weight, centroids, inertia, n_iter);
+      if (use_int64_index) {
+        _fit<double, int64_t>(
+          res, *params, X, sample_weight, centroids, inertia, n_iter);
+      } else {
+        _fit<double, int32_t>(
+          res, *params, X, sample_weight, centroids, inertia, n_iter);
+      }
     } else {
       RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d",
                 dataset.dtype.code,
@@ -282,10 +368,24 @@ extern "C" cuvsError_t cuvsKMeansPredict(cuvsResources_t res,
 {
   return cuvs::core::translate_exceptions([=] {
     auto dataset = X->dl_tensor;
+    const bool use_int64_index =
+      kmeans_predict_uses_int64_index(X, centroids, labels, params->n_clusters);
     if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) {
-        _predict<float>(res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia);
+      if (use_int64_index) {
+        _predict<float, int64_t, int64_t>(
+          res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia);
+      } else {
+        _predict<float, int32_t, int32_t>(
+          res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia);
+      }
     } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 64) {
-        _predict<double>(res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia);
+      if (use_int64_index) {
+        _predict<double, int64_t, int64_t>(
+          res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia);
+      } else {
+        _predict<double, int32_t, int32_t>(
+          res, *params, X, sample_weight, centroids, labels, normalize_weight, inertia);
+      }
     } else {
       RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d",
                 dataset.dtype.code,
@@ -301,10 +401,19 @@ extern "C" cuvsError_t cuvsKMeansClusterCost(cuvsResources_t res,
 {
   return cuvs::core::translate_exceptions([=] {
     auto dataset = X->dl_tensor;
+    const bool use_int64_index = kmeans_tensor_shapes_use_int64_index(X, centroids);
     if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) {
-      _cluster_cost<float>(res, X, centroids, cost);
+      if (use_int64_index) {
+        _cluster_cost<float, int64_t>(res, X, centroids, cost);
+      } else {
+        _cluster_cost<float, int32_t>(res, X, centroids, cost);
+      }
     } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 64) {
-      _cluster_cost<double>(res, X, centroids, cost);
+      if (use_int64_index) {
+        _cluster_cost<double, int64_t>(res, X, centroids, cost);
+      } else {
+        _cluster_cost<double, int32_t>(res, X, centroids, cost);
+      }
     } else {
       RAFT_FAIL("Unsupported dataset DLtensor dtype: %d and bits: %d",
                 dataset.dtype.code,

ci/build_standalone_c.sh

@@ -41,6 +41,9 @@ source rapids-configure-sccache
 source rapids-datetime-string
 
 rapids-pip-retry install cmake
+if [[ "${RAPIDS_CUDA_VERSION%%.*}" == "13" ]]; then
+    rapids-pip-retry install cuda-tile 'cuda-toolkit[tileiras]==13.*'
+fi
 pyenv rehash
 
 rapids-print-env

conda/environments/all_cuda-133_arch-aarch64.yaml

@@ -15,8 +15,10 @@ dependencies:
 - cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
+- cuda-tileiras
 - cuda-version=13.3
 - cupy>=14.0.1,!=14.1.0
+- cutile-python
 - cxx-compiler
 - cython>=3.2.2
 - dlpack>=0.8,<1.0

conda/environments/all_cuda-133_arch-x86_64.yaml

@@ -15,8 +15,10 @@ dependencies:
 - cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
+- cuda-tileiras
 - cuda-version=13.3
 - cupy>=14.0.1,!=14.1.0
+- cutile-python
 - cxx-compiler
 - cython>=3.2.2
 - dlpack>=0.8,<1.0

conda/environments/bench_ann_cuda-133_arch-aarch64.yaml

@@ -15,8 +15,10 @@ dependencies:
 - cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
+- cuda-tileiras
 - cuda-version=13.3
 - cupy>=14.0.1,!=14.1.0
+- cutile-python
 - cuvs==26.10.*,>=0.0.0a0
 - cxx-compiler
 - cython>=3.2.2

conda/environments/bench_ann_cuda-133_arch-x86_64.yaml

@@ -15,8 +15,10 @@ dependencies:
 - cuda-nvrtc-dev
 - cuda-nvtx-dev
 - cuda-profiler-api
+- cuda-tileiras
 - cuda-version=13.3
 - cupy>=14.0.1,!=14.1.0
+- cutile-python
 - cuvs==26.10.*,>=0.0.0a0
 - cxx-compiler
 - cython>=3.2.2

conda/recipes/libcuvs/recipe.yaml

@@ -70,6 +70,11 @@ cache:
       - cuda-version =${{ cuda_version }}
       - cmake ${{ cmake_version }}
       - ninja
+      - python
+      - if: cuda_major == "13"
+        then:
+          - cutile-python
+          - cuda-tileiras
       - ${{ stdlib("c") }}
     host:
       - libnvjitlink-dev
@@ -392,6 +397,11 @@ outputs:
         - cuda-version =${{ cuda_version }}
         - cmake ${{ cmake_version }}
         - ninja
+        - python
+        - if: cuda_major == "13"
+          then:
+            - cutile-python
+            - cuda-tileiras
         - ${{ stdlib("c") }}
       host:
         - ${{ pin_subpackage("libcuvs-headers", exact=True) }}