undo kvp, add constrains and alignments to tile export

Author: divyegala

cpp/CMakeLists.txt

@@ -980,9 +980,9 @@ if(NOT BUILD_CPU_ONLY)
     MATRIX_JSON_FILE
     "${fused_1nn_cutile_dir}/fused_1nn_cutile_matrix.json"
     FRAGMENT_TAG_FORMAT_CUBIN
-    "cuvs::distance::detail::fragment_tag_fused_1nn_cubin<cuvs::neighbors::detail::tag_@data_abbrev@, cuvs::distance::detail::metric_tag_@metric_abbrev@, cuvs::distance::detail::cutile_tile_config<@tile_m@, @tile_n@, @tile_k@>, cuvs::detail::jit_lto::@arch_tag@>"
+    "cuvs::distance::detail::fragment_tag_fused_1nn_cubin<cuvs::neighbors::detail::tag_@data_abbrev@, cuvs::distance::detail::metric_tag_@metric_abbrev@, cuvs::neighbors::detail::tag_index_@index_abbrev@, cuvs::distance::detail::cutile_tile_config<@tile_m@, @tile_n@, @tile_k@>, cuvs::detail::jit_lto::@arch_tag@>"
     FRAGMENT_TAG_FORMAT_TILEIR
-    "cuvs::distance::detail::fragment_tag_fused_1nn_tileir<cuvs::neighbors::detail::tag_@data_abbrev@, cuvs::distance::detail::metric_tag_@metric_abbrev@, cuvs::distance::detail::cutile_tile_config<@tile_m@, @tile_n@, @tile_k@>>"
+    "cuvs::distance::detail::fragment_tag_fused_1nn_tileir<cuvs::neighbors::detail::tag_@data_abbrev@, cuvs::distance::detail::metric_tag_@metric_abbrev@, cuvs::neighbors::detail::tag_index_@index_abbrev@, cuvs::distance::detail::cutile_tile_config<@tile_m@, @tile_n@, @tile_k@>>"
     FRAGMENT_TAG_HEADER_FILES
     "<cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp>"
     "<cuvs/detail/jit_lto/cutile_arch_tags.hpp>"

cpp/cmake/modules/generate_cutile_kernels.cmake

@@ -150,8 +150,22 @@ function(process_cutile_matrix_entry source_list_var)
   set(embedded_header_file "${_artifact_stem}_${register}.h")
 
   set(_python_args
-      --format "${output_format}" --data-type "${data_type}" --metric "${metric}" --tile-m
-      "${tile_m}" --tile-n "${tile_n}" --tile-k "${tile_k}" --gpu-code "${gpu_code}"
+      --format
+      "${output_format}"
+      --data-type
+      "${data_type}"
+      --metric
+      "${metric}"
+      --index-type
+      "${index_type}"
+      --tile-m
+      "${tile_m}"
+      --tile-n
+      "${tile_n}"
+      --tile-k
+      "${tile_k}"
+      --gpu-code
+      "${gpu_code}"
   )
   if(DEFINED bytecode_version AND NOT "${bytecode_version}" STREQUAL "")
     list(APPEND _python_args --bytecode-version "${bytecode_version}")

cpp/include/cuvs/detail/jit_lto/fused_distance_nn/fused_1nn_fragments.hpp

@@ -5,6 +5,8 @@
 
 #pragma once
 
+#include <cstdint>
+
 #include <cuvs/detail/jit_lto/common_fragments.hpp>
 #include <cuvs/distance/distance.hpp>
 
@@ -75,13 +77,33 @@ struct fused_1nn_data_tag<half> {
 template <typename DataT>
 using fused_1nn_data_tag_t = typename fused_1nn_data_tag<DataT>::type;
 
-template <typename DataTag, typename MetricTag, typename TileTag, typename ArchTag>
+template <typename IdxT>
+struct fused_1nn_index_tag;
+
+template <>
+struct fused_1nn_index_tag<int32_t> {
+  using type = cuvs::neighbors::detail::tag_index_i32;
+};
+
+template <>
+struct fused_1nn_index_tag<int64_t> {
+  using type = cuvs::neighbors::detail::tag_index_i64;
+};
+
+template <typename IdxT>
+using fused_1nn_index_tag_t = typename fused_1nn_index_tag<IdxT>::type;
+
+template <typename DataTag,
+          typename MetricTag,
+          typename IndexTag,
+          typename TileTag,
+          typename ArchTag>
 struct fragment_tag_fused_1nn_cubin {
   static constexpr int cc_major = ArchTag::cc_major;
   static constexpr int cc_minor = ArchTag::cc_minor;
 };
 
-template <typename DataTag, typename MetricTag, typename TileTag>
+template <typename DataTag, typename MetricTag, typename IndexTag, typename TileTag>
 struct fragment_tag_fused_1nn_tileir {};
 
 }  // namespace cuvs::distance::detail

cpp/src/cluster/detail/kmeans.cuh

@@ -682,8 +682,8 @@ void kmeans_fit(
   DataT* cur_centroids_ptr = cur_centroids_buf.data();
   DataT* new_centroids_ptr = new_centroids_buf.data();
 
-  auto minClusterAndDistance = raft::make_device_vector<raft::KeyValuePair<IndexT, DataT>, IndexT>(
-    handle, streaming_batch_size);
+  auto nearest_idx       = raft::make_device_vector<IndexT, IndexT>(handle, streaming_batch_size);
+  auto nearest_dist      = raft::make_device_vector<DataT, IndexT>(handle, streaming_batch_size);
   auto L2NormBatch       = raft::make_device_vector<DataT, IndexT>(handle, streaming_batch_size);
   auto batch_weights_buf = raft::make_device_vector<DataT, IndexT>(handle, streaming_batch_size);
   rmm::device_uvector<DataT> L2NormBuf_OR_DistBuf(0, stream);
@@ -853,8 +853,10 @@ void kmeans_fit(
         auto batch_weights_view =
           cur_batch_weights(static_cast<IndexT>(data_batch.offset()), wt_data, cur_batch_size);
 
-        auto minCAD_view = raft::make_device_vector_view<raft::KeyValuePair<IndexT, DataT>, IndexT>(
-          minClusterAndDistance.data_handle(), cur_batch_size);
+        auto nearest_idx_view =
+          raft::make_device_vector_view<IndexT, IndexT>(nearest_idx.data_handle(), cur_batch_size);
+        auto nearest_dist_view =
+          raft::make_device_vector_view<DataT, IndexT>(nearest_dist.data_handle(), cur_batch_size);
 
         if constexpr (!data_on_device) {
           if (need_compute_norms) {
@@ -883,7 +885,8 @@ void kmeans_fit(
                                      metric,
                                      iter_params.batch_samples,
                                      iter_params.batch_centroids,
-                                     minCAD_view,
+                                     nearest_idx_view,
+                                     nearest_dist_view,
                                      l2_const_view,
                                      L2NormBuf_OR_DistBuf,
                                      ws,
@@ -1071,8 +1074,7 @@ void kmeans_predict(raft::resources const& handle,
       raft::make_const_mdspan(weight.view()));
   }
 
-  auto minClusterAndDistance =
-    raft::make_device_vector<raft::KeyValuePair<IndexT, DataT>, IndexT>(handle, n_samples);
+  auto nearest_dist = raft::make_device_vector<DataT, IndexT>(handle, n_samples);
   rmm::device_uvector<DataT> L2NormBuf_OR_DistBuf(0, stream);
 
   // L2 norm of X: ||x||^2
@@ -1082,50 +1084,35 @@ void kmeans_predict(raft::resources const& handle,
     raft::linalg::norm<raft::linalg::L2Norm, raft::Apply::ALONG_ROWS>(handle, X, L2NormX.view());
   }
 
-  // computes minClusterAndDistance[0:n_samples) where  minClusterAndDistance[i]
-  // is a <key, value> pair where
-  //   'key' is index to a sample in 'centroids' (index of the nearest
-  //   centroid) and 'value' is the distance between the sample 'X[i]' and the
-  //   'centroid[key]'
   auto l2normx_view =
     raft::make_device_vector_view<const DataT, IndexT>(L2NormX.data_handle(), n_samples);
-  cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute<DataT, IndexT>(
-    handle,
-    X,
-    centroids,
-    minClusterAndDistance.view(),
-    l2normx_view,
-    L2NormBuf_OR_DistBuf,
-    pams.metric,
-    pams.batch_samples,
-    pams.batch_centroids,
-    workspace);
+  cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute<DataT, IndexT>(handle,
+                                                                             X,
+                                                                             centroids,
+                                                                             labels,
+                                                                             nearest_dist.view(),
+                                                                             l2normx_view,
+                                                                             L2NormBuf_OR_DistBuf,
+                                                                             pams.metric,
+                                                                             pams.batch_samples,
+                                                                             pams.batch_centroids,
+                                                                             workspace);
 
-  // calculate cluster cost phi_x(C)
   rmm::device_scalar<DataT> clusterCostD(stream);
-  raft::linalg::map(
-    handle,
-    minClusterAndDistance.view(),
-    [=] __device__(const raft::KeyValuePair<IndexT, DataT> kvp, DataT wt) {
-      raft::KeyValuePair<IndexT, DataT> res;
-      res.value = kvp.value * wt;
-      res.key   = kvp.key;
-      return res;
-    },
-    raft::make_const_mdspan(minClusterAndDistance.view()),
-    raft::make_const_mdspan(weight.view()));
+  raft::linalg::map(handle,
+                    nearest_dist.view(),
+                    raft::mul_op{},
+                    raft::make_const_mdspan(nearest_dist.view()),
+                    raft::make_const_mdspan(weight.view()));
 
   cuvs::cluster::kmeans::detail::computeClusterCost(
     handle,
-    minClusterAndDistance.view(),
+    nearest_dist.view(),
     workspace,
     raft::make_device_scalar_view(clusterCostD.data()),
-    raft::value_op{},
+    raft::identity_op{},
     raft::add_op{});
 
-  raft::linalg::map(
-    handle, labels, raft::key_op{}, raft::make_const_mdspan(minClusterAndDistance.view()));
-
   inertia[0] = clusterCostD.value(stream);
 }
 

cpp/src/cluster/detail/kmeans_balanced.cuh

@@ -98,58 +98,90 @@ inline std::enable_if_t<std::is_floating_point_v<MathT>> predict_core(
         raft::make_device_matrix_view<const MathT, IdxT>(centers, n_clusters, dim);
       auto X_norm_view = raft::make_device_vector_view<const MathT, IdxT>(dataset_norm, n_rows);
 
-      auto minClusterAndDistance = raft::make_device_mdarray<raft::KeyValuePair<IdxT, MathT>, IdxT>(
-        handle, mr, raft::make_extents<IdxT>(n_rows));
-
-      cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute<MathT, IdxT>(
-        handle,
-        X_view,
-        centroids_view,
-        minClusterAndDistance.view(),
-        X_norm_view,
-        L2NormBuf_OR_DistBuf,
-        params.metric,
-        0,  // batch_samples (unused for fused reduction)
-        0,  // batch_centroids (unused for fused reduction)
-        workspace);
-
-      // Copy keys to output labels
-      raft::linalg::map(handle,
-                        raft::make_const_mdspan(minClusterAndDistance.view()),
-                        raft::make_device_vector_view<LabelT, IdxT>(labels, n_rows),
-                        raft::compose_op<raft::cast_op<LabelT>, raft::key_op>());
-      break;
-    }
-    case cuvs::distance::DistanceType::InnerProduct: {
-      if (use_cutile_fused_nn<MathT, IdxT, IdxT>(handle, n_rows, n_clusters, dim)) {
-        rmm::device_uvector<MathT> L2NormBuf_OR_DistBuf(0, stream, mr);
-        rmm::device_uvector<char> workspace(0, stream, mr);
-
-        auto X_view = raft::make_device_matrix_view<const MathT, IdxT>(dataset, n_rows, dim);
-        auto centroids_view =
-          raft::make_device_matrix_view<const MathT, IdxT>(centers, n_clusters, dim);
-        auto X_norm_view = raft::make_device_vector_view<const MathT, IdxT>(dataset_norm, n_rows);
-
-        auto minClusterAndDistance =
-          raft::make_device_mdarray<raft::KeyValuePair<IdxT, MathT>, IdxT>(
-            handle, mr, raft::make_extents<IdxT>(n_rows));
+      auto nearest_dist =
+        raft::make_device_mdarray<MathT, IdxT>(handle, mr, raft::make_extents<IdxT>(n_rows));
 
+      if constexpr (std::is_same_v<LabelT, IdxT>) {
+        auto labels_view = raft::make_device_vector_view<IdxT, IdxT>(labels, n_rows);
         cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute<MathT, IdxT>(
           handle,
           X_view,
           centroids_view,
-          minClusterAndDistance.view(),
+          labels_view,
+          nearest_dist.view(),
+          X_norm_view,
+          L2NormBuf_OR_DistBuf,
+          params.metric,
+          0,  // batch_samples (unused for fused reduction)
+          0,  // batch_centroids (unused for fused reduction)
+          workspace);
+      } else {
+        auto nearest_idx =
+          raft::make_device_mdarray<IdxT, IdxT>(handle, mr, raft::make_extents<IdxT>(n_rows));
+        cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute<MathT, IdxT>(
+          handle,
+          X_view,
+          centroids_view,
+          nearest_idx.view(),
+          nearest_dist.view(),
           X_norm_view,
           L2NormBuf_OR_DistBuf,
           params.metric,
           0,
           0,
           workspace);
+        raft::copy(
+          handle, raft::make_device_vector_view<LabelT, IdxT>(labels, n_rows), nearest_idx.view());
+      }
+      break;
+    }
+    case cuvs::distance::DistanceType::InnerProduct: {
+      if (uses_fused_distance_nn(
+            use_fused<MathT, IdxT, IdxT>(handle, n_rows, n_clusters, dim, params.metric))) {
+        rmm::device_uvector<MathT> L2NormBuf_OR_DistBuf(0, stream, mr);
+        rmm::device_uvector<char> workspace(0, stream, mr);
 
-        raft::linalg::map(handle,
-                          raft::make_const_mdspan(minClusterAndDistance.view()),
-                          raft::make_device_vector_view<LabelT, IdxT>(labels, n_rows),
-                          raft::compose_op<raft::cast_op<LabelT>, raft::key_op>());
+        auto X_view = raft::make_device_matrix_view<const MathT, IdxT>(dataset, n_rows, dim);
+        auto centroids_view =
+          raft::make_device_matrix_view<const MathT, IdxT>(centers, n_clusters, dim);
+        auto X_norm_view = raft::make_device_vector_view<const MathT, IdxT>(dataset_norm, n_rows);
+
+        auto nearest_dist =
+          raft::make_device_mdarray<MathT, IdxT>(handle, mr, raft::make_extents<IdxT>(n_rows));
+
+        if constexpr (std::is_same_v<LabelT, IdxT>) {
+          auto labels_view = raft::make_device_vector_view<IdxT, IdxT>(labels, n_rows);
+          cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute<MathT, IdxT>(
+            handle,
+            X_view,
+            centroids_view,
+            labels_view,
+            nearest_dist.view(),
+            X_norm_view,
+            L2NormBuf_OR_DistBuf,
+            params.metric,
+            0,
+            0,
+            workspace);
+        } else {
+          auto nearest_idx =
+            raft::make_device_mdarray<IdxT, IdxT>(handle, mr, raft::make_extents<IdxT>(n_rows));
+          cuvs::cluster::kmeans::detail::minClusterAndDistanceCompute<MathT, IdxT>(
+            handle,
+            X_view,
+            centroids_view,
+            nearest_idx.view(),
+            nearest_dist.view(),
+            X_norm_view,
+            L2NormBuf_OR_DistBuf,
+            params.metric,
+            0,
+            0,
+            workspace);
+          raft::copy(handle,
+                     raft::make_device_vector_view<LabelT, IdxT>(labels, n_rows),
+                     nearest_idx.view());
+        }
       } else {
         rmm::device_uvector<MathT> distances(n_rows * n_clusters, stream, mr);
 
@@ -216,22 +248,19 @@ auto calc_minibatch_size(const raft::resources& handle,
   size_t mem_per_row = 0;
   switch (metric) {
     case distance::DistanceType::L2Expanded:
-    case distance::DistanceType::L2SqrtExpanded: {
-      if (use_fused<MathT, IdxT, IdxT>(handle, n_rows, n_clusters, dim)) {
-        // fusedL2NN needs a mutex and a key-value pair for each row.
-        mem_per_row += sizeof(int);
-        mem_per_row += sizeof(raft::KeyValuePair<IdxT, MathT>);
-      } else {
-        // unfused path needs a full GEMM output (distance matrix row).
-        mem_per_row += sizeof(MathT) * n_clusters;
-      }
-    } break;
+    case distance::DistanceType::L2SqrtExpanded:
     case distance::DistanceType::InnerProduct: {
-      if (use_cutile_fused_nn<MathT, IdxT, IdxT>(handle, n_rows, n_clusters, dim)) {
-        mem_per_row += sizeof(int);
-        mem_per_row += sizeof(raft::KeyValuePair<IdxT, MathT>);
-      } else {
-        mem_per_row += sizeof(MathT) * n_clusters;
+      switch (use_fused<MathT, IdxT, IdxT>(handle, n_rows, n_clusters, dim, metric)) {
+        case FusedDistancePath::FusedCutile: break;
+        case FusedDistancePath::FusedCutlass:
+          // fusedDistanceNNMinReduce CUTLASS fallback: mutex workspace + scratch KVP per row.
+          mem_per_row += sizeof(int);
+          mem_per_row += sizeof(raft::KeyValuePair<IdxT, MathT>);
+          break;
+        case FusedDistancePath::Unfused:
+          // unfused / GEMM+argmin path needs a full distance matrix row.
+          mem_per_row += sizeof(MathT) * n_clusters;
+          break;
       }
     } break;
     // Other metrics require storing a distance matrix.