[WIP] Share CUB histogram launchers

.pre-commit-config.yaml

@@ -176,6 +176,7 @@ repos:
             cpp/cmake/thirdparty/cmake/FindCUDAToolkit[.]cmake$|
             cpp/cmake/modules/FindAVX[.]cmake$|
             cpp/src/neighbors/detail/faiss_distance_utils[.]h$|
+            cpp/src/core/cub_histogram_launchers[.]hpp$|
             cpp/src/distance/detail/fused_distance_nn/gemm[.]h$|
             cpp/src/distance/detail/fused_distance_nn/epilogue[.]cuh$|
             cpp/src/distance/detail/fused_distance_nn/predicated_tile_iterator_normvec_smem[.]h$|
@@ -185,6 +186,10 @@ repos:
             cpp/src/distance/detail/fused_distance_nn/custom_epilogue_with_broadcast[.]h$|
             rust/cuvs-sys/src/bindings[.]rs$
           )
+      - id: verify-copyright
+        name: verify-copyright-cub
+        args: [--fix, --spdx, "--spdx-license-identifier=Apache-2.0 AND BSD-3-Clause"]
+        files: ^cpp/src/core/cub_histogram_launchers[.]hpp$
       - id: verify-copyright
         name: verify-copyright-cmake
         args: [--fix, --spdx, "--spdx-license-identifier=Apache-2.0 AND BSD-3-Clause"]

cpp/CMakeLists.txt

@@ -1431,8 +1431,39 @@ if(NOT BUILD_CPU_ONLY)
       "${CMAKE_CURRENT_BINARY_DIR}/src/selection/select_k_inst_data_@data_abbrev@_index_@index_abbrev@.cu"
   )
 
+  generate_inst_matrix(
+    cuvs_primitive_sources
+    MATRIX_JSON_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/core/count_labels_matrix.json"
+    INPUT_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/core/count_labels.cu.in"
+    OUTPUT_FILE_FORMAT "${CMAKE_CURRENT_BINARY_DIR}/src/core/count_labels_@instantiation_name@.cu"
+  )
+
+  generate_inst_matrix(
+    cuvs_primitive_sources
+    MATRIX_JSON_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/core/histogram_even_matrix.json"
+    INPUT_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/core/histogram_even.cu.in"
+    OUTPUT_FILE_FORMAT "${CMAKE_CURRENT_BINARY_DIR}/src/core/histogram_even_@instantiation_name@.cu"
+  )
+
+  generate_inst_matrix(
+    cuvs_primitive_sources
+    MATRIX_JSON_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/core/cub_histogram_init_matrix.json"
+    INPUT_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/core/cub_histogram_init.cu.in"
+    OUTPUT_FILE_FORMAT
+      "${CMAKE_CURRENT_BINARY_DIR}/src/core/cub_histogram_init_@instantiation_name@.cu"
+  )
+
+  generate_inst_matrix(
+    cuvs_primitive_sources
+    MATRIX_JSON_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/core/cub_histogram_sweep_matrix.json"
+    INPUT_FILE "${CMAKE_CURRENT_SOURCE_DIR}/src/core/cub_histogram_sweep.cu.in"
+    OUTPUT_FILE_FORMAT
+      "${CMAKE_CURRENT_BINARY_DIR}/src/core/cub_histogram_sweep_@instantiation_name@.cu"
+  )
+
   add_library(
     cuvs_objs OBJECT
+    ${cuvs_primitive_sources}
     src/core/cub_ptx_query.cu
     src/cluster/detail/minClusterDistanceCompute.cu
     src/cluster/agglomerative.cu

cpp/src/cluster/detail/kmeans_common.cuh

@@ -4,6 +4,7 @@
  */
 #pragma once
 
+#include "../../core/cub_histogram.hpp"
 #include "../../distance/distance.cuh"
 #include "../../distance/top_1_nn.cuh"
 #include <cstdint>
@@ -142,37 +143,49 @@ void countLabels(raft::resources const& handle,
                  IndexT n_clusters,
                  rmm::device_uvector<char>& workspace)
 {
-  cudaStream_t stream = raft::resource::get_cuda_stream(handle).get();
-
-  // CUB::DeviceHistogram requires a signed index type
-  typedef typename std::make_signed_t<IndexT> CubIndexT;
-
-  CubIndexT num_levels  = n_clusters + 1;
-  CubIndexT lower_level = 0;
-  CubIndexT upper_level = n_clusters;
-
-  size_t temp_storage_bytes = 0;
-  RAFT_CUDA_TRY(cub::DeviceHistogram::HistogramEven(nullptr,
-                                                    temp_storage_bytes,
-                                                    labels,
-                                                    count,
-                                                    num_levels,
-                                                    lower_level,
-                                                    upper_level,
-                                                    static_cast<CubIndexT>(n_samples),
-                                                    stream));
-
-  workspace.resize(temp_storage_bytes, stream);
-
-  RAFT_CUDA_TRY(cub::DeviceHistogram::HistogramEven(workspace.data(),
-                                                    temp_storage_bytes,
-                                                    labels,
-                                                    count,
-                                                    num_levels,
-                                                    lower_level,
-                                                    upper_level,
-                                                    static_cast<CubIndexT>(n_samples),
-                                                    stream));
+  if constexpr (std::is_same_v<SampleIteratorT, const uint32_t*> &&
+                (std::is_same_v<CounterT, uint32_t> ||
+                 std::is_same_v<CounterT, unsigned long long>)) {
+    using CubIndexT = std::make_signed_t<IndexT>;
+    cuvs::detail::count_labels(labels,
+                               count,
+                               static_cast<CubIndexT>(n_samples),
+                               static_cast<CubIndexT>(n_clusters),
+                               workspace,
+                               raft::resource::get_cuda_stream(handle).get());
+  } else {
+    cudaStream_t stream = raft::resource::get_cuda_stream(handle).get();
+
+    // CUB::DeviceHistogram requires a signed index type
+    typedef typename std::make_signed_t<IndexT> CubIndexT;
+
+    CubIndexT num_levels  = n_clusters + 1;
+    CubIndexT lower_level = 0;
+    CubIndexT upper_level = n_clusters;
+
+    size_t temp_storage_bytes = 0;
+    RAFT_CUDA_TRY(cuvs::detail::histogram_even(nullptr,
+                                               temp_storage_bytes,
+                                               labels,
+                                               count,
+                                               num_levels,
+                                               lower_level,
+                                               upper_level,
+                                               static_cast<CubIndexT>(n_samples),
+                                               stream));
+
+    workspace.resize(temp_storage_bytes, stream);
+
+    RAFT_CUDA_TRY(cuvs::detail::histogram_even(workspace.data(),
+                                               temp_storage_bytes,
+                                               labels,
+                                               count,
+                                               num_levels,
+                                               lower_level,
+                                               upper_level,
+                                               static_cast<CubIndexT>(n_samples),
+                                               stream));
+  }
 }
 
 /**

cpp/src/core/count_labels.cu.in

@@ -0,0 +1,51 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <core/cub_histogram.hpp>
+#include <core/cub_ptx_query.hpp>
+#include <cub/device/device_histogram.cuh>
+#include <cuvs/core/cuda_fp16.hpp>
+#include <raft/util/cudart_utils.hpp>
+
+namespace cuvs::detail {
+
+template <typename CounterT, typename IndexT>
+void count_labels(const uint32_t* labels,
+                  CounterT* counts,
+                  IndexT n_samples,
+                  IndexT n_clusters,
+                  rmm::device_uvector<char>& workspace,
+                  cudaStream_t stream)
+{
+  IndexT num_levels         = n_clusters + 1;
+  IndexT lower_level        = 0;
+  IndexT upper_level        = n_clusters;
+  size_t temp_storage_bytes = 0;
+  RAFT_CUDA_TRY(cuvs::detail::histogram_even(nullptr,
+                                             temp_storage_bytes,
+                                             labels,
+                                             counts,
+                                             num_levels,
+                                             lower_level,
+                                             upper_level,
+                                             n_samples,
+                                             stream));
+  workspace.resize(temp_storage_bytes, stream);
+  RAFT_CUDA_TRY(cuvs::detail::histogram_even(workspace.data(),
+                                             temp_storage_bytes,
+                                             labels,
+                                             counts,
+                                             num_levels,
+                                             lower_level,
+                                             upper_level,
+                                             n_samples,
+                                             stream));
+}
+
+// clang-format off
+template @return_type@ @qualified_function@@template_arguments@(@parameter_types@)@function_qualifiers@;
+// clang-format on
+
+}  // namespace cuvs::detail

cpp/src/core/count_labels_matrix.json

@@ -0,0 +1,30 @@
+[
+  {
+    "return_type": "void",
+    "qualified_function": "count_labels",
+    "function_qualifiers": "",
+    "namespace_context": "cuvs::detail",
+    "_variants": [
+      {
+        "template_arguments": "<uint32_t, int>",
+        "parameter_types": "const uint32_t*, uint32_t*, int, int, rmm::device_uvector<char>&, cudaStream_t",
+        "instantiation_name": "u32_int"
+      },
+      {
+        "template_arguments": "<uint32_t, int64_t>",
+        "parameter_types": "const uint32_t*, uint32_t*, int64_t, int64_t, rmm::device_uvector<char>&, cudaStream_t",
+        "instantiation_name": "u32_i64"
+      },
+      {
+        "template_arguments": "<unsigned long long, int>",
+        "parameter_types": "const uint32_t*, unsigned long long*, int, int, rmm::device_uvector<char>&, cudaStream_t",
+        "instantiation_name": "u64_int"
+      },
+      {
+        "template_arguments": "<unsigned long long, int64_t>",
+        "parameter_types": "const uint32_t*, unsigned long long*, int64_t, int64_t, rmm::device_uvector<char>&, cudaStream_t",
+        "instantiation_name": "u64_i64"
+      }
+    ]
+  }
+]

cpp/src/core/cub_histogram.hpp

@@ -0,0 +1,36 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <cstdint>
+#include <cuda_runtime_api.h>
+#include <rmm/device_uvector.hpp>
+
+namespace cuvs::detail {
+
+/** CUB label histograms shared by balanced k-means instantiations. */
+template <typename CounterT, typename IndexT>
+void count_labels(const uint32_t* labels,
+                  CounterT* counts,
+                  IndexT n_samples,
+                  IndexT n_clusters,
+                  rmm::device_uvector<char>& workspace,
+                  cudaStream_t stream);
+
+}  // namespace cuvs::detail
+
+namespace cuvs::detail {
+template <typename SampleIt, typename CounterT, typename LevelT, typename NumItems>
+cudaError_t histogram_even(void* storage,
+                           size_t& bytes,
+                           SampleIt samples,
+                           CounterT* counts,
+                           int levels,
+                           LevelT lower,
+                           LevelT upper,
+                           NumItems count,
+                           cudaStream_t stream);
+}  // namespace cuvs::detail

cpp/src/core/cub_histogram_init.cu.in

@@ -0,0 +1,19 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <core/cub_histogram_launchers.hpp>
+#include <core/cub_ptx_query.hpp>
+
+namespace cuvs::detail::histogram_launch {
+template <typename Policy, int Channels, typename Counter, typename Offset>
+CUB_RUNTIME_FUNCTION init_pointer_t<Policy, Channels, Counter, Offset> init()
+{
+  return &cub::detail::histogram::DeviceHistogramInitKernel<Policy, Channels, Counter, Offset>;
+}
+}  // namespace cuvs::detail::histogram_launch
+
+// clang-format off
+template @return_type@ @qualified_function@@template_arguments@();
+// clang-format on

cpp/src/core/cub_histogram_init_matrix.json

@@ -0,0 +1,62 @@
+[
+  {
+    "return_type": "cuvs::detail::histogram_launch::init_pointer_t<cub::detail::histogram::policy_selector_from_types<unsigned int, unsigned long long, 1, 1, true>, 1, unsigned long long, int>",
+    "qualified_function": "cuvs::detail::histogram_launch::init",
+    "template_arguments": "<cub::detail::histogram::policy_selector_from_types<unsigned int, unsigned long long, 1, 1, true>, 1, unsigned long long, int>",
+    "instantiation_name": "u32_count_u64_offset_i32"
+  },
+  {
+    "return_type": "cuvs::detail::histogram_launch::init_pointer_t<cub::detail::histogram::policy_selector_from_types<int, float, 1, 1, true>, 1, float, int>",
+    "qualified_function": "cuvs::detail::histogram_launch::init",
+    "template_arguments": "<cub::detail::histogram::policy_selector_from_types<int, float, 1, 1, true>, 1, float, int>",
+    "instantiation_name": "i32_count_f32_offset_i32"
+  },
+  {
+    "return_type": "cuvs::detail::histogram_launch::init_pointer_t<cub::detail::histogram::policy_selector_from_types<long, double, 1, 1, true>, 1, double, long>",
+    "qualified_function": "cuvs::detail::histogram_launch::init",
+    "template_arguments": "<cub::detail::histogram::policy_selector_from_types<long, double, 1, 1, true>, 1, double, long>",
+    "instantiation_name": "i64_count_f64_offset_i64"
+  },
+  {
+    "return_type": "cuvs::detail::histogram_launch::init_pointer_t<cub::detail::histogram::policy_selector_from_types<long, double, 1, 1, true>, 1, double, int>",
+    "qualified_function": "cuvs::detail::histogram_launch::init",
+    "template_arguments": "<cub::detail::histogram::policy_selector_from_types<long, double, 1, 1, true>, 1, double, int>",
+    "instantiation_name": "i64_count_f64_offset_i32"
+  },
+  {
+    "return_type": "cuvs::detail::histogram_launch::init_pointer_t<cub::detail::histogram::policy_selector_from_types<long, float, 1, 1, true>, 1, float, long>",
+    "qualified_function": "cuvs::detail::histogram_launch::init",
+    "template_arguments": "<cub::detail::histogram::policy_selector_from_types<long, float, 1, 1, true>, 1, float, long>",
+    "instantiation_name": "i64_count_f32_offset_i64"
+  },
+  {
+    "return_type": "cuvs::detail::histogram_launch::init_pointer_t<cub::detail::histogram::policy_selector_from_types<long, float, 1, 1, true>, 1, float, int>",
+    "qualified_function": "cuvs::detail::histogram_launch::init",
+    "template_arguments": "<cub::detail::histogram::policy_selector_from_types<long, float, 1, 1, true>, 1, float, int>",
+    "instantiation_name": "i64_count_f32_offset_i32"
+  },
+  {
+    "return_type": "cuvs::detail::histogram_launch::init_pointer_t<cub::detail::histogram::policy_selector_from_types<unsigned int, unsigned int, 1, 1, true>, 1, unsigned int, int>",
+    "qualified_function": "cuvs::detail::histogram_launch::init",
+    "template_arguments": "<cub::detail::histogram::policy_selector_from_types<unsigned int, unsigned int, 1, 1, true>, 1, unsigned int, int>",
+    "instantiation_name": "u32_count_u32_offset_i32"
+  },
+  {
+    "return_type": "cuvs::detail::histogram_launch::init_pointer_t<cub::detail::histogram::policy_selector_from_types<int, double, 1, 1, true>, 1, double, int>",
+    "qualified_function": "cuvs::detail::histogram_launch::init",
+    "template_arguments": "<cub::detail::histogram::policy_selector_from_types<int, double, 1, 1, true>, 1, double, int>",
+    "instantiation_name": "i32_count_f64_offset_i32"
+  },
+  {
+    "return_type": "cuvs::detail::histogram_launch::init_pointer_t<cub::detail::histogram::policy_selector_from_types<int, int, 1, 1, true>, 1, int, int>",
+    "qualified_function": "cuvs::detail::histogram_launch::init",
+    "template_arguments": "<cub::detail::histogram::policy_selector_from_types<int, int, 1, 1, true>, 1, int, int>",
+    "instantiation_name": "i32_count_i32_offset_i32"
+  },
+  {
+    "return_type": "cuvs::detail::histogram_launch::init_pointer_t<cub::detail::histogram::policy_selector_from_types<unsigned int, unsigned int, 1, 1, true>, 1, unsigned int, long>",
+    "qualified_function": "cuvs::detail::histogram_launch::init",
+    "template_arguments": "<cub::detail::histogram::policy_selector_from_types<unsigned int, unsigned int, 1, 1, true>, 1, unsigned int, long>",
+    "instantiation_name": "u32_count_u32_offset_i64"
+  }
+]