Remove JIT+LTO fragment database (rapidsai#1927)

Rather than register each fragment in a runtime class with a string key, "register" them with the linker using template specialization. This solves a number of problems:

1. It simplifies the code by removing the `FragmentDatabase` class.
2. It addresses rapidsai#1909 (comment) by bypassing the issue entirely. There is no longer a need to build the fragment name string at runtime.
3. For clients that use the `cuvs_static` static library, it allows the linker to pick and choose which fragment symbols it needs rather than including all of them with every client just in case any of them are needed.
4. Since there is no longer a need for `$<WHOLE_ARCHIVE:...>` linkage, there is no need for the `cuvs_jit_lto_kernels` target at all, thus simplifying the CMake code too.

Authors:
  - Kyle Edwards (https://github.com/KyleFromNVIDIA)

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

URL: rapidsai#1927

Author: KyleFromNVIDIA

cpp/CMakeLists.txt

@@ -357,6 +357,7 @@ if(NOT BUILD_CPU_ONLY)
 
   set(JIT_LTO_TARGET_ARCHITECTURE "")
   set(JIT_LTO_COMPILATION OFF)
+  set(jit_lto_files)
   if(CMAKE_CUDA_COMPILER_VERSION VERSION_GREATER_EQUAL 13.0)
     set(JIT_LTO_TARGET_ARCHITECTURE "75-real")
     set(JIT_LTO_COMPILATION ON)
@@ -434,28 +435,16 @@ if(NOT BUILD_CPU_ONLY)
     )
     endblock()
 
-    add_library(
-      cuvs_jit_lto_kernels STATIC
-      ${interleaved_scan_files}
-      ${metric_files}
-      ${filter_files}
-      ${post_lambda_files}
-      src/detail/jit_lto/AlgorithmLauncher.cpp
-      src/detail/jit_lto/AlgorithmPlanner.cpp
-      src/detail/jit_lto/FragmentDatabase.cpp
-      src/detail/jit_lto/FragmentEntry.cpp
-      src/detail/jit_lto/nvjitlink_checker.cpp
+    set(jit_lto_files
+        ${interleaved_scan_files}
+        ${metric_files}
+        ${filter_files}
+        ${post_lambda_files}
+        src/detail/jit_lto/AlgorithmLauncher.cpp
+        src/detail/jit_lto/AlgorithmPlanner.cpp
+        src/detail/jit_lto/FragmentEntry.cpp
+        src/detail/jit_lto/nvjitlink_checker.cpp
     )
-    set_target_properties(
-      cuvs_jit_lto_kernels PROPERTIES POSITION_INDEPENDENT_CODE ON CXX_STANDARD 20
-    )
-    target_include_directories(
-      cuvs_jit_lto_kernels
-      PRIVATE "${CMAKE_CURRENT_SOURCE_DIR}/include" "${CMAKE_CURRENT_SOURCE_DIR}/src"
-              "${CMAKE_CURRENT_SOURCE_DIR}/../c/include"
-    )
-    target_link_libraries(cuvs_jit_lto_kernels PRIVATE raft::raft)
-    add_library(cuvs::cuvs_jit_lto_kernels ALIAS cuvs_jit_lto_kernels)
   endif()
 
   add_library(
@@ -666,6 +655,7 @@ if(NOT BUILD_CPU_ONLY)
     src/stats/silhouette_score.cu
     src/stats/trustworthiness_score.cu
     ${CUVS_MG_ALGOS}
+    ${jit_lto_files}
   )
 
   set_target_properties(
@@ -777,12 +767,8 @@ if(NOT BUILD_CPU_ONLY)
              $<BUILD_LOCAL_INTERFACE:$<TARGET_NAME_IF_EXISTS:NCCL::NCCL>>
              $<BUILD_LOCAL_INTERFACE:$<TARGET_NAME_IF_EXISTS:hnswlib::hnswlib>>
              $<$<BOOL:${CUVS_NVTX}>:CUDA::nvtx3>
-      PRIVATE
-        $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX>
-        $<COMPILE_ONLY:nvidia::cutlass::cutlass>
-        $<COMPILE_ONLY:cuco::cuco>
-        $<$<BOOL:${JIT_LTO_COMPILATION}>:CUDA::nvJitLink>
-        $<$<BOOL:${JIT_LTO_COMPILATION}>:$<LINK_LIBRARY:WHOLE_ARCHIVE,cuvs::cuvs_jit_lto_kernels>>
+      PRIVATE $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX> $<COMPILE_ONLY:nvidia::cutlass::cutlass>
+              $<COMPILE_ONLY:cuco::cuco> $<$<BOOL:${JIT_LTO_COMPILATION}>:CUDA::nvJitLink>
     )
 
     # ensure CUDA symbols aren't relocated to the middle of the debug build binaries
@@ -839,13 +825,11 @@ SECTIONS
              ${CUVS_CTK_MATH_DEPENDENCIES}
              $<TARGET_NAME_IF_EXISTS:NCCL::NCCL> # needs to be public for DT_NEEDED
              $<BUILD_LOCAL_INTERFACE:$<TARGET_NAME_IF_EXISTS:hnswlib::hnswlib>> # header only
-      PRIVATE
-        $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX>
-        $<$<BOOL:${JIT_LTO_COMPILATION}>:CUDA::nvJitLink>
-        $<$<BOOL:${CUVS_NVTX}>:CUDA::nvtx3>
-        $<COMPILE_ONLY:nvidia::cutlass::cutlass>
-        $<COMPILE_ONLY:cuco::cuco>
-        $<$<BOOL:${JIT_LTO_COMPILATION}>:$<LINK_LIBRARY:WHOLE_ARCHIVE,cuvs::cuvs_jit_lto_kernels>>
+      PRIVATE $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX>
+              $<$<BOOL:${JIT_LTO_COMPILATION}>:CUDA::nvJitLink>
+              $<$<BOOL:${CUVS_NVTX}>:CUDA::nvtx3>
+              $<COMPILE_ONLY:nvidia::cutlass::cutlass>
+              $<COMPILE_ONLY:cuco::cuco>
     )
   endif()
 
@@ -886,11 +870,9 @@ target_compile_definitions(cuvs::cuvs INTERFACE $<$<BOOL:${CUVS_NVTX}>:NVTX_ENAB
   include(GNUInstallDirs)
   include(CPack)
 
-  set(target_names cuvs cuvs_static cuvs_jit_lto_kernels cuvs_cpp_headers cuvs_c)
-  set(component_names cuvs_shared cuvs_static cuvs_static cuvs_cpp_headers c_api)
-  set(export_names cuvs-shared-exports cuvs-static-exports cuvs-static-exports
-                   cuvs-cpp-headers-exports cuvs-c-exports
-  )
+  set(target_names cuvs cuvs_static cuvs_cpp_headers cuvs_c)
+  set(component_names cuvs_shared cuvs_static cuvs_cpp_headers c_api)
+  set(export_names cuvs-shared-exports cuvs-static-exports cuvs-cpp-headers-exports cuvs-c-exports)
   foreach(target component export IN ZIP_LISTS target_names component_names export_names)
     if(TARGET ${target})
       install(

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

@@ -14,21 +14,22 @@
 struct FragmentEntry;
 
 struct AlgorithmPlanner {
-  AlgorithmPlanner(std::string fragment_key, std::string entrypoint)
-    : fragment_key(std::move(fragment_key)), entrypoint(std::move(entrypoint))
-  {
-  }
+  AlgorithmPlanner(std::string entrypoint) : entrypoint(std::move(entrypoint)) {}
 
   std::shared_ptr<AlgorithmLauncher> get_launcher();
 
-  std::string fragment_key;
   std::string entrypoint;
-  std::vector<std::string> device_functions;
-  std::vector<FragmentEntry*> fragments;
+  std::vector<const FragmentEntry*> fragments;
+
+  void add_fragment(const FragmentEntry& fragment);
+
+  template <typename FragmentT>
+  void add_fragment()
+  {
+    add_fragment(FragmentT{});
+  }
 
  private:
-  void add_entrypoint();
-  void add_device_functions();
-  std::string get_device_functions_key() const;
+  std::string get_fragments_key() const;
   std::shared_ptr<AlgorithmLauncher> build();
 };

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

@@ -12,21 +12,27 @@
 
 #include <nvJitLink.h>
 
+#include "nvjitlink_checker.hpp"
+
 struct FragmentEntry {
-  FragmentEntry(std::string const& key);
+  virtual bool add_to(nvJitLinkHandle& handle) const = 0;
 
-  bool operator==(const FragmentEntry& rhs) const { return compute_key == rhs.compute_key; }
+  virtual const char* get_key() const = 0;
+};
 
-  virtual bool add_to(nvJitLinkHandle& handle) const = 0;
+struct FatbinFragmentEntry : FragmentEntry {
+  virtual const uint8_t* get_data() const = 0;
+
+  virtual size_t get_length() const = 0;
 
-  std::string compute_key{};
+  bool add_to(nvJitLinkHandle& handle) const override final;
 };
 
-struct FatbinFragmentEntry final : FragmentEntry {
-  FatbinFragmentEntry(std::string const& key, unsigned char const* view, std::size_t size);
+template <typename FragmentT>
+struct StaticFatbinFragmentEntry : FatbinFragmentEntry {
+  const uint8_t* get_data() const override final { return FragmentT::data; }
 
-  virtual bool add_to(nvJitLinkHandle& handle) const;
+  size_t get_length() const override final { return FragmentT::length; }
 
-  std::size_t data_size          = 0;
-  unsigned char const* data_view = nullptr;
+  const char* get_key() const override final { return typeid(FragmentT).name(); }
 };

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

@@ -0,0 +1,52 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <cuvs/detail/jit_lto/FragmentEntry.hpp>
+
+namespace cuvs::neighbors::ivf_flat::detail {
+
+template <typename DataTag,
+          typename AccTag,
+          typename IdxTag,
+          int Capacity,
+          int Veclen,
+          bool Ascending,
+          bool ComputeNorm>
+struct InterleavedScanFragmentEntry final
+  : StaticFatbinFragmentEntry<InterleavedScanFragmentEntry<DataTag,
+                                                           AccTag,
+                                                           IdxTag,
+                                                           Capacity,
+                                                           Veclen,
+                                                           Ascending,
+                                                           ComputeNorm>> {
+  static const uint8_t* const data;
+  static const size_t length;
+};
+
+template <int Veclen, typename DataTag, typename AccTag, typename MetricTag>
+struct MetricFragmentEntry final
+  : StaticFatbinFragmentEntry<MetricFragmentEntry<Veclen, DataTag, AccTag, MetricTag>> {
+  static const uint8_t* const data;
+  static const size_t length;
+};
+
+template <typename IvfSampleFilterTag>
+struct FilterFragmentEntry final
+  : StaticFatbinFragmentEntry<FilterFragmentEntry<IvfSampleFilterTag>> {
+  static const uint8_t* const data;
+  static const size_t length;
+};
+
+template <typename PostLambdaTag>
+struct PostLambdaFragmentEntry final
+  : StaticFatbinFragmentEntry<PostLambdaFragmentEntry<PostLambdaTag>> {
+  static const uint8_t* const data;
+  static const size_t length;
+};
+
+}  // namespace cuvs::neighbors::ivf_flat::detail

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

@@ -70,10 +70,7 @@ template <typename IdxTag, typename FilterImplTag>
 struct tag_filter {};
 
 // Tag types for distance metrics with full template info
-template <int Veclen, typename TTag, typename AccTTag>
 struct tag_metric_euclidean {};
-
-template <int Veclen, typename TTag, typename AccTTag>
 struct tag_metric_inner_product {};
 
 // Tag types for post-processing

cpp/include/cuvs/detail/jit_lto/ivf_flat/interleaved_scan_fragments.hpp

@@ -3,8 +3,6 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-#include "nvjitlink_checker.hpp"
-
 #include <chrono>
 #include <iterator>
 #include <memory>
@@ -14,51 +12,41 @@
 #include <vector>
 
 #include <cuvs/detail/jit_lto/AlgorithmPlanner.hpp>
-#include <cuvs/detail/jit_lto/FragmentDatabase.hpp>
+#include <cuvs/detail/jit_lto/FragmentEntry.hpp>
+#include <cuvs/detail/jit_lto/nvjitlink_checker.hpp>
 
 #include "cuda_runtime.h"
 #include "nvJitLink.h"
 
 #include <raft/core/logger.hpp>
 #include <raft/util/cuda_rt_essentials.hpp>
 
-void AlgorithmPlanner::add_entrypoint()
+void AlgorithmPlanner::add_fragment(const FragmentEntry& fragment)
 {
-  auto entrypoint_fragment = fragment_database().get_fragment(this->fragment_key);
-  this->fragments.push_back(entrypoint_fragment);
-}
-
-void AlgorithmPlanner::add_device_functions()
-{
-  for (const auto& device_function_key : this->device_functions) {
-    auto device_function_fragment = fragment_database().get_fragment(device_function_key);
-    this->fragments.push_back(device_function_fragment);
-  }
+  fragments.push_back(&fragment);
 }
 
-std::string AlgorithmPlanner::get_device_functions_key() const
+std::string AlgorithmPlanner::get_fragments_key() const
 {
   std::string key = "";
-  for (const auto& device_function : this->device_functions) {
-    key += device_function;
+  for (const auto* fragment : this->fragments) {
+    key += fragment->get_key();
   }
   return key;
 }
 
 std::shared_ptr<AlgorithmLauncher> AlgorithmPlanner::get_launcher()
 {
   auto& launchers = get_cached_launchers();
-  auto launch_key = this->fragment_key + this->get_device_functions_key();
+  auto launch_key = this->get_fragments_key();
 
   static std::mutex cache_mutex;
   std::lock_guard<std::mutex> lock(cache_mutex);
   if (launchers.count(launch_key) == 0) {
-    add_entrypoint();
-    add_device_functions();
     std::string log_message =
-      "JIT compiling launcher for kernel: " + this->fragment_key + " and device functions: ";
-    for (const auto& device_function : this->device_functions) {
-      log_message += device_function + ",";
+      "JIT compiling launcher for kernel: " + this->entrypoint + " and device functions: ";
+    for (const auto* fragment : this->fragments) {
+      log_message += std::string{fragment->get_key()} + ",";
     }
     log_message.pop_back();
     RAFT_LOG_INFO("%s", log_message.c_str());