[WIP] Share RAFT keyed row-reduction kernels

cpp/CMakeLists.txt

@@ -1431,8 +1431,38 @@ if(NOT BUILD_CPU_ONLY)
       "${CMAKE_CURRENT_BINARY_DIR}/src/selection/select_k_inst_data_@data_abbrev@_index_@index_abbrev@.cu"
   )
 
+  generate_string_matrix(
+    reduce_rows_by_key_detail_extern_instantiations
+    MATRIX_JSON_FILE
+    "${CMAKE_CURRENT_SOURCE_DIR}/src/core/reduce_rows_by_key_detail_matrix.json"
+    ITEM_FORMAT
+    "extern template @return_type@ @qualified_function@@template_arguments@(@parameter_types@)@function_qualifiers@@semicolon@"
+    GLUE
+    "\n"
+  )
+  generate_string_matrix(
+    reduce_rows_by_key_helpers_extern_instantiations
+    MATRIX_JSON_FILE
+    "${CMAKE_CURRENT_SOURCE_DIR}/src/core/reduce_rows_by_key_helpers_matrix.json"
+    ITEM_FORMAT
+    "extern template @return_type@ @qualified_function@@template_arguments@(@parameter_types@)@function_qualifiers@@semicolon@"
+    GLUE
+    "\n"
+  )
+  foreach(primitive IN ITEMS src/core/reduce_rows_by_key src/core/reduce_rows_by_key_detail
+                             src/core/reduce_rows_by_key_helpers
+  )
+    generate_inst_matrix(
+      cuvs_primitive_sources
+      MATRIX_JSON_FILE "${CMAKE_CURRENT_SOURCE_DIR}/${primitive}_matrix.json"
+      INPUT_FILE "${CMAKE_CURRENT_SOURCE_DIR}/${primitive}.cu.in"
+      OUTPUT_FILE_FORMAT "${CMAKE_CURRENT_BINARY_DIR}/${primitive}_@instantiation_name@.cu"
+    )
+  endforeach()
+
   add_library(
     cuvs_objs OBJECT
+    ${cuvs_primitive_sources}
     src/cluster/detail/minClusterDistanceCompute.cu
     src/cluster/agglomerative.cu
     ${gmm_inst_files}

cpp/src/cluster/detail/kmeans_balanced.cuh

@@ -5,6 +5,7 @@
 
 #pragma once
 
+#include "../../core/keyed_row_reduction_launchers.hpp"
 #include "kmeans_common.cuh"
 #include <cuvs/cluster/kmeans.hpp>
 
@@ -415,7 +416,7 @@ void calc_centers_and_sizes(const raft::resources& handle,
   if (is_packed_binary) {
     if constexpr (std::is_same_v<T, uint8_t>) {
       auto decoded_dataset_iter = make_bitwise_expanded_iterator<MathT, IdxT>(dataset);
-      raft::linalg::reduce_rows_by_key(decoded_dataset_iter,
+      cuvs::detail::reduce_rows_by_key(decoded_dataset_iter,
                                        centers_dim,
                                        labels,
                                        nullptr,
@@ -431,7 +432,7 @@ void calc_centers_and_sizes(const raft::resources& handle,
   }
   // Apply mapping only when the data and math types are different.
   else if constexpr (std::is_same_v<T, MathT>) {
-    raft::linalg::reduce_rows_by_key(dataset,
+    cuvs::detail::reduce_rows_by_key(dataset,
                                      dim,
                                      labels,
                                      nullptr,
@@ -444,7 +445,7 @@ void calc_centers_and_sizes(const raft::resources& handle,
   } else {
     // todo(lsugy): use iterator from KV output of fusedL2NN
     thrust::transform_iterator<MappingOpT, const T*> mapping_itr(dataset, mapping_op);
-    raft::linalg::reduce_rows_by_key(mapping_itr,
+    cuvs::detail::reduce_rows_by_key(mapping_itr,
                                      dim,
                                      labels,
                                      nullptr,

cpp/src/cluster/detail/kmeans_common.cuh

@@ -4,6 +4,7 @@
  */
 #pragma once
 
+#include "../../core/keyed_row_reduction_launchers.hpp"
 #include "../../distance/distance.cuh"
 #include "../../distance/top_1_nn.cuh"
 #include <cstdint>
@@ -637,7 +638,7 @@ void compute_centroid_adjustments(
 
   workspace.resize(n_samples, stream);
 
-  raft::linalg::reduce_rows_by_key(X.data_handle(),
+  cuvs::detail::reduce_rows_by_key(X.data_handle(),
                                    X.extent(1),
                                    cluster_labels,
                                    sample_weights.data_handle(),

cpp/src/core/keyed_row_reduction_launchers.hpp

@@ -0,0 +1,53 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <optional>
+#include <raft/core/device_mdspan.hpp>
+#include <raft/core/resources.hpp>
+
+namespace cuvs::detail {
+using raft::row_major;
+template <typename DataIteratorT,
+          typename KeysIteratorT,
+          typename WeightT,
+          typename SumsT,
+          typename IdxT>
+void reduce_rows_by_key(const DataIteratorT d_A,
+                        IdxT lda,
+                        const KeysIteratorT d_keys,
+                        const WeightT* d_weights,
+                        char* d_keys_char,
+                        IdxT nrows,
+                        IdxT ncols,
+                        IdxT nkeys,
+                        SumsT* d_sums,
+                        cudaStream_t stream,
+                        bool reset_sums = true);
+
+template <typename DataIteratorT, typename KeysIteratorT, typename SumsT, typename IdxT>
+void reduce_rows_by_key(const DataIteratorT d_A,
+                        IdxT lda,
+                        const KeysIteratorT d_keys,
+                        char* d_keys_char,
+                        IdxT nrows,
+                        IdxT ncols,
+                        IdxT nkeys,
+                        SumsT* d_sums,
+                        cudaStream_t stream,
+                        bool reset_sums = true);
+
+template <typename ElementType, typename KeyType, typename WeightType, typename IndexType>
+void reduce_rows_by_key(
+  raft::resources const& handle,
+  raft::device_matrix_view<const ElementType, IndexType, raft::row_major> d_A,
+  raft::device_vector_view<const KeyType, IndexType> d_keys,
+  raft::device_matrix_view<ElementType, IndexType, raft::row_major> d_sums,
+  IndexType n_unique_keys,
+  raft::device_vector_view<char, IndexType> d_keys_char,
+  std::optional<raft::device_vector_view<const WeightType, IndexType>> d_weights = std::nullopt,
+  bool reset_sums                                                                = true);
+}  // namespace cuvs::detail

cpp/src/core/reduce_rows_by_key.cu.in

@@ -0,0 +1,81 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <distance/top_1_nn.cuh>
+#include <neighbors/detail/ann_utils.cuh>
+
+#include <core/keyed_row_reduction_launchers.hpp>
+
+#include <cuvs/core/cuda_fp16.hpp>
+#include <raft/linalg/reduce_rows_by_key.cuh>
+#include <thrust/iterator/counting_iterator.h>
+#include <thrust/iterator/transform_iterator.h>
+
+namespace cuvs::detail {
+
+template <typename DataIteratorT,
+          typename KeysIteratorT,
+          typename WeightT,
+          typename SumsT,
+          typename IdxT>
+void reduce_rows_by_key(const DataIteratorT d_A,
+                        IdxT lda,
+                        const KeysIteratorT d_keys,
+                        const WeightT* d_weights,
+                        char* d_keys_char,
+                        IdxT nrows,
+                        IdxT ncols,
+                        IdxT nkeys,
+                        SumsT* d_sums,
+                        cudaStream_t stream,
+                        bool reset_sums)
+
+{
+  return raft::linalg::reduce_rows_by_key(
+    d_A, lda, d_keys, d_weights, d_keys_char, nrows, ncols, nkeys, d_sums, stream, reset_sums);
+}
+
+template <typename DataIteratorT, typename KeysIteratorT, typename SumsT, typename IdxT>
+void reduce_rows_by_key(const DataIteratorT d_A,
+                        IdxT lda,
+                        const KeysIteratorT d_keys,
+                        char* d_keys_char,
+                        IdxT nrows,
+                        IdxT ncols,
+                        IdxT nkeys,
+                        SumsT* d_sums,
+                        cudaStream_t stream,
+                        bool reset_sums)
+
+{
+  return raft::linalg::reduce_rows_by_key(
+    d_A, lda, d_keys, d_keys_char, nrows, ncols, nkeys, d_sums, stream, reset_sums);
+}
+
+template <typename ElementType, typename KeyType, typename WeightType, typename IndexType>
+void reduce_rows_by_key(
+  raft::resources const& handle,
+  raft::device_matrix_view<const ElementType, IndexType, raft::row_major> d_A,
+  raft::device_vector_view<const KeyType, IndexType> d_keys,
+  raft::device_matrix_view<ElementType, IndexType, raft::row_major> d_sums,
+  IndexType n_unique_keys,
+  raft::device_vector_view<char, IndexType> d_keys_char,
+  std::optional<raft::device_vector_view<const WeightType, IndexType>> d_weights,
+  bool reset_sums)
+
+{
+  return raft::linalg::reduce_rows_by_key(
+    handle, d_A, d_keys, d_sums, n_unique_keys, d_keys_char, d_weights, reset_sums);
+}
+
+}  // namespace cuvs::detail
+
+// clang-format off
+@reduce_rows_by_key_helpers_extern_instantiations@
+
+@reduce_rows_by_key_detail_extern_instantiations@
+
+template @return_type@ @qualified_function@@template_arguments@(@parameter_types@)@function_qualifiers@;
+// clang-format on

cpp/src/core/reduce_rows_by_key_detail.cu.in

@@ -0,0 +1,20 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <distance/top_1_nn.cuh>
+#include <neighbors/detail/ann_utils.cuh>
+
+#include <core/keyed_row_reduction_launchers.hpp>
+
+#include <cuvs/core/cuda_fp16.hpp>
+#include <raft/linalg/reduce_rows_by_key.cuh>
+#include <thrust/iterator/counting_iterator.h>
+#include <thrust/iterator/transform_iterator.h>
+
+// clang-format off
+@reduce_rows_by_key_helpers_extern_instantiations@
+
+template @return_type@ @qualified_function@@template_arguments@(@parameter_types@)@function_qualifiers@;
+// clang-format on

cpp/src/core/reduce_rows_by_key_detail_matrix.json

@@ -0,0 +1,24 @@
+[
+  {
+    "return_type": "void",
+    "qualified_function": "raft::linalg::detail::reduce_rows_by_key",
+    "function_qualifiers": "",
+    "_variants": [
+      {
+        "template_arguments": "< thrust::transform_iterator<cuvs::spatial::knn::detail::utils::mapping<float>, __half const*, thrust::use_default, thrust::use_default>, unsigned int const*, float, float, long>",
+        "parameter_types": "thrust::transform_iterator<cuvs::spatial::knn::detail::utils::mapping<float>, __half const*, thrust::use_default, thrust::use_default>, long, unsigned int const*, float const*, char*, long, long, long, float*, CUstream_st*, bool",
+        "instantiation_name": "raft_linalg_detail_reduce_rows_by_key_raw_transform_mapping_f32_f16_c_ptr_use_default_use_default_u32_c_ptr_f32_f32_long"
+      },
+      {
+        "template_arguments": "< thrust::transform_iterator<cuvs::spatial::knn::detail::utils::mapping<float>, signed char const*, thrust::use_default, thrust::use_default>, unsigned int const*, float, float, long>",
+        "parameter_types": "thrust::transform_iterator<cuvs::spatial::knn::detail::utils::mapping<float>, signed char const*, thrust::use_default, thrust::use_default>, long, unsigned int const*, float const*, char*, long, long, long, float*, CUstream_st*, bool",
+        "instantiation_name": "raft_linalg_detail_reduce_rows_by_key_raw_transform_mapping_f32_i8_c_ptr_use_default_use_default_u32_c_ptr_f32_f32_long"
+      },
+      {
+        "template_arguments": "< thrust::transform_iterator<cuvs::spatial::knn::detail::utils::mapping<float>, unsigned char const*, thrust::use_default, thrust::use_default>, unsigned int const*, float, float, long>",
+        "parameter_types": "thrust::transform_iterator<cuvs::spatial::knn::detail::utils::mapping<float>, unsigned char const*, thrust::use_default, thrust::use_default>, long, unsigned int const*, float const*, char*, long, long, long, float*, CUstream_st*, bool",
+        "instantiation_name": "raft_linalg_detail_reduce_rows_by_key_raw_transform_mapping_f32_u8_c_ptr_use_default_use_default_u32_c_ptr_f32_f32_long"
+      }
+    ]
+  }
+]

cpp/src/core/reduce_rows_by_key_helpers.cu.in

@@ -0,0 +1,10 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <raft/linalg/reduce_rows_by_key.cuh>
+
+// clang-format off
+template @return_type@ @qualified_function@@template_arguments@(@parameter_types@)@function_qualifiers@;
+// clang-format on

cpp/src/core/reduce_rows_by_key_helpers_matrix.json

@@ -0,0 +1,58 @@
+[
+  {
+    "return_type": "void",
+    "qualified_function": "raft::linalg::detail::sum_rows_by_key_small_nkeys",
+    "template_arguments": "<const float*, float, float, int>",
+    "parameter_types": "const float*, int, const char*, const float*, int, int, int, float*, cudaStream_t",
+    "function_qualifiers": "",
+    "instantiation_name": "sum_f32_i32"
+  },
+  {
+    "return_type": "void",
+    "qualified_function": "raft::linalg::detail::sum_rows_by_key_small_nkeys",
+    "template_arguments": "<const float*, float, float, int64_t>",
+    "parameter_types": "const float*, int64_t, const char*, const float*, int64_t, int64_t, int64_t, float*, cudaStream_t",
+    "function_qualifiers": "",
+    "instantiation_name": "sum_f32_i64"
+  },
+  {
+    "return_type": "void",
+    "qualified_function": "raft::linalg::detail::sum_rows_by_key_small_nkeys",
+    "template_arguments": "<const double*, double, double, int>",
+    "parameter_types": "const double*, int, const char*, const double*, int, int, int, double*, cudaStream_t",
+    "function_qualifiers": "",
+    "instantiation_name": "sum_f64_i32"
+  },
+  {
+    "return_type": "void",
+    "qualified_function": "raft::linalg::detail::sum_rows_by_key_small_nkeys",
+    "template_arguments": "<const double*, double, double, int64_t>",
+    "parameter_types": "const double*, int64_t, const char*, const double*, int64_t, int64_t, int64_t, double*, cudaStream_t",
+    "function_qualifiers": "",
+    "instantiation_name": "sum_f64_i64"
+  },
+  {
+    "return_type": "void",
+    "qualified_function": "raft::linalg::detail::convert_array",
+    "template_arguments": "<char*, int*>",
+    "parameter_types": "char*, int*, int, cudaStream_t",
+    "function_qualifiers": "",
+    "instantiation_name": "convert_i32"
+  },
+  {
+    "return_type": "void",
+    "qualified_function": "raft::linalg::detail::convert_array",
+    "template_arguments": "<char*, int64_t*>",
+    "parameter_types": "char*, int64_t*, int, cudaStream_t",
+    "function_qualifiers": "",
+    "instantiation_name": "convert_i64"
+  },
+  {
+    "return_type": "void",
+    "qualified_function": "raft::linalg::detail::convert_array",
+    "template_arguments": "<char*, const uint32_t*>",
+    "parameter_types": "char*, const uint32_t*, int, cudaStream_t",
+    "function_qualifiers": "",
+    "instantiation_name": "convert_u32"
+  }
+]