Use Specific CCCL Includes (#1806)

This is needed downstream because cub injects an EmptyKernel symbol into every TU that includes either cub/cub.cuh or cub/device/* includes, and this causes an issue for cubins created using JIT-LTO.

This PR also fixes 2 bugs in JIT kernels: 
1. Filters out `cub::EmptyKernel` from JIT TUs specifically because they only expect 1 kernel but are finding 2 instead
2. Keeps a reference to `cudaLibrary_t` that loads the JIT kernel and unloading it at destruction

Authors:
  - Divye Gala (https://github.com/divyegala)

Approvers:
  - Dante Gama Dessavre (https://github.com/dantegd)

URL: #1806

Author: divyegala

cpp/include/cuvs/detail/jit_lto/AlgorithmLauncher.hpp

@@ -15,9 +15,17 @@
 #include <memory>
 
 struct AlgorithmLauncher {
-  AlgorithmLauncher() = default;
+  AlgorithmLauncher() : kernel{nullptr}, library{nullptr} {}
 
-  AlgorithmLauncher(cudaKernel_t k);
+  AlgorithmLauncher(cudaKernel_t k, cudaLibrary_t lib);
+
+  ~AlgorithmLauncher();
+
+  AlgorithmLauncher(const AlgorithmLauncher&)            = delete;
+  AlgorithmLauncher& operator=(const AlgorithmLauncher&) = delete;
+
+  AlgorithmLauncher(AlgorithmLauncher&& other) noexcept;
+  AlgorithmLauncher& operator=(AlgorithmLauncher&& other) noexcept;
 
   template <typename... Args>
   void dispatch(cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, Args&&... args)
@@ -31,6 +39,7 @@ struct AlgorithmLauncher {
  private:
   void call(cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** args);
   cudaKernel_t kernel;
+  cudaLibrary_t library;
 };
 
 std::unordered_map<std::string, std::shared_ptr<AlgorithmLauncher>>& get_cached_launchers();

cpp/src/cluster/detail/kmeans_common.cuh

@@ -31,7 +31,10 @@
 #include <rmm/device_scalar.hpp>
 #include <rmm/device_uvector.hpp>
 
-#include <cub/cub.cuh>
+#include <cub/device/device_histogram.cuh>
+#include <cub/device/device_reduce.cuh>
+#include <cub/device/device_select.cuh>
+#include <cub/iterator/arg_index_input_iterator.cuh>
 #include <cuda.h>
 #include <thrust/fill.h>
 #include <thrust/for_each.h>

cpp/src/detail/jit_lto/AlgorithmLauncher.cu

@@ -7,7 +7,32 @@
 
 #include <raft/util/cuda_rt_essentials.hpp>
 
-AlgorithmLauncher::AlgorithmLauncher(cudaKernel_t k) : kernel{k} {}
+AlgorithmLauncher::AlgorithmLauncher(cudaKernel_t k, cudaLibrary_t lib) : kernel{k}, library{lib} {}
+
+AlgorithmLauncher::~AlgorithmLauncher()
+{
+  if (library != nullptr) { (void)cudaLibraryUnload(library); }
+}
+
+AlgorithmLauncher::AlgorithmLauncher(AlgorithmLauncher&& other) noexcept
+  : kernel{other.kernel}, library{other.library}
+{
+  other.kernel  = nullptr;
+  other.library = nullptr;
+}
+
+AlgorithmLauncher& AlgorithmLauncher::operator=(AlgorithmLauncher&& other) noexcept
+{
+  if (this != &other) {
+    // Unload current library if it exists
+    if (library != nullptr) { cudaLibraryUnload(library); }
+    kernel        = other.kernel;
+    library       = other.library;
+    other.kernel  = nullptr;
+    other.library = nullptr;
+  }
+  return *this;
+}
 
 void AlgorithmLauncher::call(
   cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** kernel_args)

cpp/src/detail/jit_lto/AlgorithmPlanner.cu

@@ -106,15 +106,44 @@ std::shared_ptr<AlgorithmLauncher> AlgorithmPlanner::build()
   RAFT_EXPECTS(result == NVJITLINK_SUCCESS, "nvJitLinkDestroy failed");
 
   // cubin is linked, so now load it
-  // NOTE: cudaLibrary_t does not need to be freed explicitly
   cudaLibrary_t library;
   RAFT_CUDA_TRY(
     cudaLibraryLoadData(&library, cubin.get(), nullptr, nullptr, 0, nullptr, nullptr, 0));
 
-  constexpr unsigned int count = 1;
+  unsigned int kernel_count = 0;
+  RAFT_CUDA_TRY(cudaLibraryGetKernelCount(&kernel_count, library));
+
   // NOTE: cudaKernel_t does not need to be freed explicitly
-  std::unique_ptr<cudaKernel_t[]> kernels{new cudaKernel_t[count]};
-  RAFT_CUDA_TRY(cudaLibraryEnumerateKernels(kernels.get(), count, library));
+  std::unique_ptr<cudaKernel_t[]> kernels{new cudaKernel_t[kernel_count]};
+  RAFT_CUDA_TRY(cudaLibraryEnumerateKernels(kernels.get(), kernel_count, library));
+
+  // Filter out EmptyKernel by checking kernel names using cudaFuncGetName
+  const char* empty_kernel_name = "_ZN3cub6detail11EmptyKernelIvEEvv";
+  std::vector<cudaKernel_t> valid_kernels;
+  valid_kernels.reserve(kernel_count);
+
+  for (unsigned int i = 0; i < kernel_count; ++i) {
+    // cudaFuncGetName can be used with cudaKernel_t by casting to void*
+    const void* func_ptr  = reinterpret_cast<const void*>(kernels[i]);
+    const char* func_name = nullptr;
+    RAFT_CUDA_TRY(cudaFuncGetName(&func_name, func_ptr));
+
+    bool is_empty_kernel = false;
+    if (func_name != nullptr) {
+      std::string kernel_name(func_name);
+      // Check if this is EmptyKernel
+      if (kernel_name.find(empty_kernel_name) != std::string::npos ||
+          kernel_name == empty_kernel_name) {
+        is_empty_kernel = true;
+      }
+    }
+
+    // Only keep the kernel if it's not EmptyKernel
+    if (!is_empty_kernel) { valid_kernels.push_back(kernels[i]); }
+  }
+
+  RAFT_EXPECTS(
+    valid_kernels.size() == 1, "Expected 1 valid JIT kernel, got %zu", valid_kernels.size());
 
-  return std::make_shared<AlgorithmLauncher>(kernels.release()[0]);
+  return std::make_shared<AlgorithmLauncher>(valid_kernels[0], library);
 }

cpp/src/distance/detail/sparse/coo_spmv_kernel.cuh

@@ -11,7 +11,7 @@
 #include <cub/block/block_load.cuh>
 #include <cub/block/block_radix_sort.cuh>
 #include <cub/block/block_store.cuh>
-#include <cub/cub.cuh>
+#include <cub/warp/warp_reduce.cuh>
 
 namespace cuvs {
 namespace distance {

cpp/src/distance/detail/sparse/utils.cuh

@@ -1,13 +1,13 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #pragma once
 
 #include <raft/core/math.hpp>
 
-#include <cub/cub.cuh>
+#include <cub/warp/warp_reduce.cuh>
 #include <cuda_fp16.h>
 #include <cuda_pipeline.h>
 

cpp/src/distance/fused_distance_nn-inl.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -14,7 +14,7 @@
 #include <raft/linalg/contractions.cuh>
 #include <raft/util/cuda_utils.cuh>
 
-#include <cub/cub.cuh>
+#include <cub/util_type.cuh>
 
 #include <stdint.h>
 

cpp/src/neighbors/detail/cagra/topk_for_cagra/topk_core.cuh

@@ -7,7 +7,7 @@
 #include "../utils.hpp"
 #include "topk.h"
 
-#include <cub/cub.cuh>
+#include <cub/block/block_scan.cuh>
 
 #include <raft/core/detail/macros.hpp>
 #include <raft/core/error.hpp>

cpp/src/neighbors/detail/fused_l2_knn.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2021-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
@@ -8,7 +8,7 @@
 #include <raft/linalg/norm.cuh>
 #include <raft/util/cuda_utils.cuh>
 
-#include <cub/cub.cuh>
+#include <cub/util_type.cuh>
 
 #include <limits>
 

cpp/src/neighbors/detail/knn_utils.cuh

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -9,7 +9,6 @@
 
 #include <cuvs/distance/distance.hpp>
 
-#include <cub/cub.cuh>
 #include <cuda_fp16.h>
 #include <cuda_pipeline.h>