JIT LTO Cagra Search

conda/recipes/libcuvs/recipe.yaml

@@ -399,13 +399,13 @@ outputs:
         - librmm =${{ minor_version }}
         - nccl ${{ nccl_version }}
         - cuda-cudart-dev
+        - cuda-nvrtc-dev
         - cuda-profiler-api
         - libcublas-dev
         - libcurand-dev
         - libcusolver-dev
         - libcusparse-dev
         - libnvjitlink-dev
-        - cuda-nvrtc-dev
       run:
         - ${{ pin_subpackage("libcuvs-headers", exact=True) }}
         - ${{ pin_subpackage("libcuvs", exact=True) }}

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

@@ -37,10 +37,24 @@ struct AlgorithmLauncher {
     this->call(stream, grid, block, shared_mem, kernel_args);
   }
 
+  template <typename FuncT, typename... Args>
+  void dispatch_cooperative(
+    cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, Args&&... args)
+  {
+    static_assert(
+      std::is_same_v<FuncT, void(std::remove_reference_t<Args>...)>,
+      "dispatch_cooperative() argument types do not match the kernel function signature FuncT");
+
+    void* kernel_args[] = {const_cast<void*>(static_cast<void const*>(&args))...};
+    this->call_cooperative(stream, grid, block, shared_mem, kernel_args);
+  }
+
   cudaKernel_t get_kernel() { return this->kernel; }
 
  private:
   void call(cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** args);
+  void call_cooperative(
+    cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** args);
   cudaKernel_t kernel;
   cudaLibrary_t library;
 };

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

@@ -44,6 +44,12 @@ struct AlgorithmPlanner {
     add_fragment(std::make_unique<StaticFatbinFragmentEntry<FragmentTag>>());
   }
 
+ protected:
+  /** Extra link-time option strings passed to nvJitLink. Base build()
+   *  always passes "-lto" and "-arch=sm_XX" first; derived planners may append here in their
+   *  constructor body. */
+  std::vector<std::string> linktime_extra_options;
+
  private:
   std::string get_fragments_key() const;
   std::shared_ptr<AlgorithmLauncher> build();

cpp/include/cuvs/detail/jit_lto/cagra/cagra_fragments.hpp

@@ -0,0 +1,93 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <cstdint>
+
+namespace cuvs::neighbors::cagra::detail {
+
+struct tag_dist_f {};
+struct tag_metric_l2 {};
+struct tag_metric_inner_product {};
+struct tag_metric_cosine {};
+struct tag_metric_hamming {};
+struct tag_codebook_none {};
+struct tag_codebook_half {};
+struct tag_metric_l1 {};
+struct tag_norm_noop {};
+struct tag_norm_cosine {};
+
+/// Multi-kernel planners that do not link `sample_filter` into the JIT link (e.g.
+/// `random_pickup`). Real filters use `cuvs::neighbors::detail::tag_filter_*` on
+/// `CagraPlannerBase`.
+struct tag_cagra_jit_sample_filter_link_absent {};
+
+template <typename DataTag,
+          typename IndexTag,
+          typename DistanceTag,
+          typename QueryTag,
+          typename CodebookTag,
+          uint32_t TeamSize,
+          uint32_t DatasetBlockDim,
+          uint32_t PqBits,
+          uint32_t PqLen>
+struct fragment_tag_setup_workspace {};
+
+template <typename DataTag,
+          typename IndexTag,
+          typename DistanceTag,
+          typename QueryTag,
+          typename CodebookTag,
+          uint32_t TeamSize,
+          uint32_t DatasetBlockDim,
+          uint32_t PqBits,
+          uint32_t PqLen>
+struct fragment_tag_compute_distance {};
+
+template <typename QueryTag, typename DistanceTag, typename MetricTag>
+struct fragment_tag_dist_op {};
+
+template <typename DataTag,
+          typename IndexTag,
+          typename DistanceTag,
+          typename QueryTag,
+          uint32_t TeamSize,
+          uint32_t DatasetBlockDim,
+          typename NormTag>
+struct fragment_tag_apply_normalization_standard {};
+
+template <typename DataTag,
+          typename SourceIndexTag,
+          typename IndexTag,
+          typename DistanceTag,
+          bool TopkByBitonicSort,
+          bool BitonicSortAndMergeMultiWarps>
+struct fragment_tag_search_single_cta {};
+
+template <typename DataTag,
+          typename SourceIndexTag,
+          typename IndexTag,
+          typename DistanceTag,
+          bool TopkByBitonicSort,
+          bool BitonicSortAndMergeMultiWarps>
+struct fragment_tag_search_single_cta_p {};
+
+template <typename DataTag, typename SourceIndexTag, typename IndexTag, typename DistanceTag>
+struct fragment_tag_search_multi_cta {};
+
+template <typename DataTag, typename IndexTag, typename DistanceTag>
+struct fragment_tag_random_pickup {};
+
+template <typename DataTag, typename IndexTag, typename DistanceTag, typename SourceIndexTag>
+struct fragment_tag_compute_distance_to_child_nodes {};
+
+template <typename IndexTag, typename DistanceTag, typename SourceIndexTag>
+struct fragment_tag_apply_filter_kernel {};
+
+template <typename BitsetTag, typename SourceIndexTag, typename FilterTag>
+struct fragment_tag_sample_filter {};
+
+}  // namespace cuvs::neighbors::cagra::detail

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

@@ -7,11 +7,16 @@
 
 namespace cuvs::neighbors::detail {
 
+struct tag_f {};
+struct tag_h {};
+struct tag_i8 {};
+struct tag_u8 {};
 struct tag_filter_none {};
 struct tag_filter_bitset {};
 
 struct tag_bitset_u32 {};
 
+struct tag_index_u32 {};
 struct tag_index_i64 {};
 
 template <typename BitsetTag, typename IndexTag, typename FilterTag>

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

@@ -7,19 +7,13 @@
 
 namespace cuvs::neighbors::ivf_flat::detail {
 
-// Tag types for data types
-struct tag_f {};
-struct tag_h {};
-struct tag_i8 {};
-struct tag_u8 {};
-
 // Tag types for accumulator types
 struct tag_acc_f {};
 struct tag_acc_h {};
 struct tag_acc_i32 {};
 struct tag_acc_u32 {};
 
-// Tag types for distance metrics with full template info
+// Tag types for distance metrics
 struct tag_metric_euclidean {};
 struct tag_metric_inner_product {};
 struct tag_metric_custom_udf {};

cpp/src/detail/jit_lto/AlgorithmLauncher.cpp

@@ -24,7 +24,6 @@ AlgorithmLauncher::AlgorithmLauncher(AlgorithmLauncher&& other) noexcept
 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;
@@ -47,3 +46,21 @@ void AlgorithmLauncher::call(
 
   RAFT_CUDA_TRY(cudaLaunchKernelExC(&config, kernel, kernel_args));
 }
+
+void AlgorithmLauncher::call_cooperative(
+  cudaStream_t stream, dim3 grid, dim3 block, std::size_t shared_mem, void** kernel_args)
+{
+  cudaLaunchAttribute attribute[1];
+  attribute[0].id              = cudaLaunchAttributeCooperative;
+  attribute[0].val.cooperative = 1;
+
+  cudaLaunchConfig_t config{};
+  config.gridDim          = grid;
+  config.blockDim         = block;
+  config.stream           = stream;
+  config.dynamicSmemBytes = shared_mem;
+  config.numAttrs         = 1;
+  config.attrs            = attribute;
+
+  RAFT_CUDA_TRY(cudaLaunchKernelExC(&config, kernel, kernel_args));
+}

cpp/src/detail/jit_lto/AlgorithmPlanner.cpp

@@ -73,8 +73,14 @@ std::shared_ptr<AlgorithmLauncher> AlgorithmPlanner::build()
 
   // Load the generated LTO IR and link them together
   nvJitLinkHandle handle;
-  const char* lopts[] = {"-lto", archs.c_str()};
-  auto result         = nvJitLinkCreate(&handle, 2, lopts);
+  std::vector<const char*> lopts;
+  lopts.reserve(2 + linktime_extra_options.size());
+  lopts.push_back("-lto");
+  lopts.push_back(archs.c_str());
+  for (auto const& opt : linktime_extra_options) {
+    lopts.push_back(opt.c_str());
+  }
+  auto result = nvJitLinkCreate(&handle, static_cast<unsigned int>(lopts.size()), lopts.data());
   check_nvjitlink_result(handle, result);
 
   for (const auto& frag : this->fragments) {

cpp/src/neighbors/cagra.cuh

@@ -9,7 +9,7 @@
 #include "detail/cagra/cagra_build.cuh"
 #include "detail/cagra/cagra_merge.cuh"
 #include "detail/cagra/cagra_search.cuh"
-#include "detail/cagra/graph_core.cuh"
+#include "detail/cagra/jit_lto_kernels/graph_core.cuh"
 
 #include "detail/ann_utils.cuh"
 #include <raft/core/device_mdspan.hpp>

cpp/src/neighbors/detail/cagra/cagra_build.cuh

@@ -6,7 +6,7 @@
 
 #include "../../../core/nvtx.hpp"
 #include "../../../preprocessing/quantize/vpq_build-ext.cuh"
-#include "graph_core.cuh"
+#include "jit_lto_kernels/graph_core.cuh"
 
 #include <raft/core/copy.cuh>
 #include <raft/core/device_mdarray.hpp>

cpp/src/neighbors/detail/cagra/cagra_search.cuh

@@ -32,6 +32,7 @@
 #include <raft/linalg/norm.cuh>
 #include <raft/linalg/reduce.cuh>
 
+// All includes are done before opening namespace to avoid nested namespace issues
 namespace cuvs::neighbors::cagra::detail {
 
 template <typename DataT,

cpp/src/neighbors/detail/cagra/compute_distance.hpp

@@ -1,11 +1,12 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION.
  * SPDX-License-Identifier: Apache-2.0
  */
 #pragma once
 
-#include "device_common.hpp"
-#include "hashmap.hpp"
+#include "../neighbors_device_intrinsics.cuh"
+#include "jit_lto_kernels/device_memory_ops.hpp"
+#include "jit_lto_kernels/hashmap.hpp"
 #include "utils.hpp"
 
 #include <cuvs/distance/distance.hpp>
@@ -137,33 +138,20 @@ struct alignas(device::LOAD_128BIT_T) dataset_descriptor_base_t {
   };
   static_assert(sizeof(smem_and_team_size_t) == sizeof(uint32_t));
 
-  using setup_workspace_type  = const base_type*(const base_type*, void*, const DATA_T*, uint32_t);
-  using compute_distance_type = DISTANCE_T(const args_t, const INDEX_T);
-
   args_t args;
 
-  /** Copy the descriptor and the query into shared memory and do any other work, such as
-   * initializing the codebook. */
-  setup_workspace_type* setup_workspace_impl;
-  /** Compute the distance from the query vector (stored in the smem_workspace) and a dataset vector
-   * given by the dataset_index. */
-  compute_distance_type* compute_distance_impl;
   /** A placeholder for an implementation-specific pointer. */
   void* extra_ptr3;
   smem_and_team_size_t smem_and_team_size;
 
   /** Number of records in the database. */
   INDEX_T size;
 
-  RAFT_INLINE_FUNCTION dataset_descriptor_base_t(setup_workspace_type* setup_workspace_impl,
-                                                 compute_distance_type* compute_distance_impl,
-                                                 INDEX_T size,
+  RAFT_INLINE_FUNCTION dataset_descriptor_base_t(INDEX_T size,
                                                  uint32_t dim,
                                                  uint32_t team_size_bitshift,
                                                  uint32_t smem_ws_size_in_bytes)
-    : setup_workspace_impl(setup_workspace_impl),
-      compute_distance_impl(compute_distance_impl),
-      size(size),
+    : size(size),
       smem_and_team_size(smem_ws_size_in_bytes, team_size_bitshift),
       args{nullptr, nullptr, 0, dim, 0, 0}
   {
@@ -191,20 +179,6 @@ struct alignas(device::LOAD_128BIT_T) dataset_descriptor_base_t {
   {
     return smem_and_team_size.team_size();
   }
-
-  RAFT_DEVICE_INLINE_FUNCTION auto setup_workspace(void* smem_ptr,
-                                                   const DATA_T* queries_ptr,
-                                                   uint32_t query_id) const -> const base_type*
-  {
-    return setup_workspace_impl(this, smem_ptr, queries_ptr, query_id);
-  }
-
-  RAFT_DEVICE_INLINE_FUNCTION auto compute_distance(INDEX_T dataset_index, bool valid) const
-    -> DISTANCE_T
-  {
-    auto per_thread_distances = valid ? compute_distance_impl(args.load(), dataset_index) : 0;
-    return device::team_sum(per_thread_distances, team_size_bitshift_from_smem());
-  }
 };
 
 /**
@@ -227,6 +201,14 @@ struct dataset_descriptor_host {
   uint32_t smem_ws_size_in_bytes = 0;
   uint32_t team_size             = 0;
 
+  // JIT LTO metadata - stored when descriptor is created
+  cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded;
+  uint32_t dataset_block_dim          = 0;
+  bool is_vpq                         = false;
+  uint32_t pq_bits                    = 0;
+  uint32_t pq_len                     = 0;
+  // Codebook type is determined by DataT for VPQ (always half for now)
+
   struct state {
     using ready_t = std::tuple<dev_descriptor_t*, rmm::cuda_stream_view>;
     using init_f =
@@ -270,10 +252,21 @@ struct dataset_descriptor_host {
   };
 
   template <typename DescriptorImpl, typename InitF>
-  dataset_descriptor_host(const DescriptorImpl& dd_host, InitF init)
+  dataset_descriptor_host(const DescriptorImpl& dd_host,
+                          InitF init,
+                          cuvs::distance::DistanceType metric_val,
+                          uint32_t dataset_block_dim_val,
+                          bool is_vpq_val      = false,
+                          uint32_t pq_bits_val = 0,
+                          uint32_t pq_len_val  = 0)
     : value_{std::make_shared<state>(init, sizeof(DescriptorImpl))},
       smem_ws_size_in_bytes{dd_host.smem_ws_size_in_bytes()},
-      team_size{dd_host.team_size()}
+      team_size{dd_host.team_size()},
+      metric{metric_val},
+      dataset_block_dim{dataset_block_dim_val},
+      is_vpq{is_vpq_val},
+      pq_bits{pq_bits_val},
+      pq_len{pq_len_val}
   {
   }