Add JIT-LTO based filter UDF support for CAGRA (#2132)

This PR adds a first version of low-level JIT-LTO filter UDF support for CAGRA search in cuVS C++.

Users can provide CUDA device source for a predicate with this ABI:

```cpp
__device__ bool cuvs_filter_udf(uint32_t query_id,
                                source_index_t source_id,
                                void* filter_data);
```

The predicate returns `true` to allow a candidate and `false` to reject it. `filter_data` is not result data; it is an optional opaque device-accessible context pointer that lets the predicate read user metadata such as tenant ids, timestamps, language masks, ACL bitmaps, or query-specific thresholds. Results are still written to the normal CAGRA `neighbors` and `distances` outputs.

This gives CAGRA a path to support runtime metadata predicates without ahead-of-time template explosion, while keeping existing `none` and `bitset` filters on their static JIT-LTO fragment paths.

## What Changed

- Added `cuvs::neighbors::filtering::udf_filter`.
- Added `FilterType::UDF`.
- Added CAGRA dispatch for `udf_filter`, including `filtering_rate` fallback behavior.
- Added a dynamic JIT-LTO sample-filter fragment generator for UDF source.
- Reworked CAGRA JIT sample-filter payload plumbing from bitset-only payloads to a generic payload that supports:
  - no filter
  - bitset filter
  - UDF filter context pointer
  - query id offsets for batched/multi-kernel paths
- Updated single-CTA, multi-CTA, and multi-kernel CAGRA JIT paths to call the linked sample filter uniformly.
- Added docs for UDF behavior, caveats, and non-goals.
- Added a standalone C++ example showing three metadata-style UDFs.

## About `filter_data`

`filter_data` is the mechanism for passing runtime metadata into the device predicate. It is optional: simple predicates can ignore it or use `nullptr`.

For example, this UDF needs no context:

```cpp
__device__ bool cuvs_filter_udf(uint32_t, source_index_t source_id, void*)
{
  return source_id >= 704;
}
```

For metadata filters, callers pass a device pointer to a user-defined context struct:

```cpp
struct tenant_filter_context {
  const uint32_t* row_tenant_ids;
  const uint32_t* query_tenant_ids;
};
```

Then the UDF casts `filter_data` back to that type:

```cpp
__device__ bool tenant_filter(uint32_t query_id,
                              source_index_t source_id,
                              void* filter_data)
{
  auto* ctx = static_cast<const tenant_filter_context*>(filter_data);
  return ctx->row_tenant_ids[source_id] == ctx->query_tenant_ids[query_id];
}
```

The pointer and anything it points to must be device-accessible and remain valid for the duration of the search. Future typed wrappers or expression builders could hide this cast, but the internal ABI can stay stable as:

```cpp
bool predicate(uint32_t query_id, source_index_t source_id, void* filter_data);
```

## Example

The standalone example builds one CAGRA index and runs three different UDF predicates over the same metadata context.

First, the caller defines a host-side context type. This same layout is repeated in the UDF source string so device code knows how to interpret `filter_data`:

```cpp
struct metadata_filter_context {
  const uint32_t* row_tenant_ids;
  const int64_t* row_timestamps;
  const uint32_t* row_language_ids;
  const uint64_t* row_acl_masks;

  const uint32_t* query_tenant_ids;
  const int64_t* query_min_timestamps;
  const uint64_t* query_allowed_language_masks;
  const uint64_t* query_permission_masks;
};
```

The UDF source contains the device predicates. Each predicate receives the same `filter_data` pointer and casts it to `metadata_filter_context`:

```cpp
std::string source = R"cpp(
struct metadata_filter_context {
  const uint32_t* row_tenant_ids;
  const int64_t* row_timestamps;
  const uint32_t* row_language_ids;
  const uint64_t* row_acl_masks;

  const uint32_t* query_tenant_ids;
  const int64_t* query_min_timestamps;
  const uint64_t* query_allowed_language_masks;
  const uint64_t* query_permission_masks;
};

__device__ bool tenant_filter(uint32_t query_id,
                              source_index_t source_id,
                              void* filter_data)
{
  auto* ctx = static_cast<const metadata_filter_context*>(filter_data);
  return ctx->row_tenant_ids[source_id] == ctx->query_tenant_ids[query_id];
}

__device__ bool timestamp_filter(uint32_t query_id,
                                 source_index_t source_id,
                                 void* filter_data)
{
  auto* ctx = static_cast<const metadata_filter_context*>(filter_data);
  return ctx->row_timestamps[source_id] >= ctx->query_min_timestamps[query_id];
}

__device__ bool language_acl_filter(uint32_t query_id,
                                    source_index_t source_id,
                                    void* filter_data)
{
  auto* ctx = static_cast<const metadata_filter_context*>(filter_data);
  const auto language_bit = uint64_t{1} << ctx->row_language_ids[source_id];
  const bool language_ok =
    (ctx->query_allowed_language_masks[query_id] & language_bit) != 0;
  const bool acl_ok =
    (ctx->row_acl_masks[source_id] & ctx->query_permission_masks[query_id]) != 0;
  return language_ok && acl_ok;
}
)cpp";
```

The caller owns the metadata arrays. They must be copied to device memory before search:

```cpp
auto row_tenant_ids_device = raft::make_device_vector<uint32_t, int64_t>(res, n_rows);
auto row_timestamps_device = raft::make_device_vector<int64_t, int64_t>(res, n_rows);
auto row_language_ids_device = raft::make_device_vector<uint32_t, int64_t>(res, n_rows);
auto row_acl_masks_device = raft::make_device_vector<uint64_t, int64_t>(res, n_rows);

auto query_tenant_ids_device = raft::make_device_vector<uint32_t, int64_t>(res, n_queries);
auto query_min_timestamps_device = raft::make_device_vector<int64_t, int64_t>(res, n_queries);
auto query_allowed_language_masks_device =
  raft::make_device_vector<uint64_t, int64_t>(res, n_queries);
auto query_permission_masks_device =
  raft::make_device_vector<uint64_t, int64_t>(res, n_queries);

// Copy host metadata into these device arrays before launching search.
```

Then the caller builds one device-resident context struct whose fields point at those device arrays:

```cpp
auto context_device =
  raft::make_device_vector<metadata_filter_context, int64_t>(res, 1);

metadata_filter_context host_context{
  row_tenant_ids_device.data_handle(),
  row_timestamps_device.data_handle(),
  row_language_ids_device.data_handle(),
  row_acl_masks_device.data_handle(),
  query_tenant_ids_device.data_handle(),
  query_min_timestamps_device.data_handle(),
  query_allowed_language_masks_device.data_handle(),
  query_permission_masks_device.data_handle()
};

// Copy the struct itself to device memory. The pointer to this device struct
// is what we pass as filter_data.
raft::copy(context_device.data_handle(),
           &host_context,
           1,
           raft::resource::get_cuda_stream(res));
raft::resource::sync_stream(res);
```

Finally, each `udf_filter` selects which device function to link by name. All three filters reuse the same source and the same `filter_data` context:

```cpp
auto tenant_filter = cuvs::neighbors::filtering::udf_filter(
  source,
  context_device.data_handle(),  // filter_data
  0.75f,                         // estimated rejected fraction
  "tenant-filter-v1",            // cache key
  "tenant_filter");              // device function name

auto timestamp_filter = cuvs::neighbors::filtering::udf_filter(
  source,
  context_device.data_handle(),
  0.50f,
  "timestamp-filter-v1",
  "timestamp_filter");

auto language_acl_filter = cuvs::neighbors::filtering::udf_filter(
  source,
  context_device.data_handle(),
  0.875f,
  "language-acl-filter-v1",
  "language_acl_filter");
```

Example output:

```text
Building CAGRA index
tenant_filter first query neighbors: 3364 3488 260 772 3292 3344 2052 660
timestamp_filter first query neighbors: 4070 3942 3364 3629 3797 3488 3292 3495
language_acl_filter first query neighbors: 3488 3344 2480 2520 2304 2176 1360 400
All CAGRA filter UDF examples produced valid filtered neighbors.
```

## Validation

Focused CAGRA UDF test passes across all CAGRA search algorithms:

```text
NEIGHBORS_ANN_CAGRA_FILTER_UDF_TEST
```

Coverage includes:

- accept-all UDF matches no-filter results
- reject-all UDF returns no valid dataset row ids
- high `filtering_rate` UDF returns only accepted rows
- invalid source throws during compile/search setup
- repeated same-cache-key searches match
- UDF threshold predicate matches equivalent bitset filter
- query-specific tenant metadata works across SINGLE_CTA, MULTI_CTA, and MULTI_KERNEL

Broader CAGRA regression set also passed:

```text
NEIGHBORS_ANN_CAGRA_(FILTER_UDF_TEST|FLOAT_UINT32|INT8_UINT32|UINT8_UINT32|HALF_UINT32|TEST_BUGS)
```

## Notes

This PR intentionally keeps the v1 API low-level: users provide CUDA source plus an optional `void*` device context. Typed wrappers, expression builders, or write-once host/device ergonomics can be layered on later without changing the internal ABI.

Filter UDFs are candidate-validity predicates only. They do not control CAGRA graph traversal, distance computation, PQ/VPQ internals, or result selection.

## Benchmarks

Performance benchmarking is still TODO for this PR.

The expected regression risk for existing functionality is low because `none` and `bitset` filters still use static JIT-LTO fragments, and the UDF dynamic fragment path is only used for `udf_filter`. However, the CAGRA JIT kernel payload/signature plumbing changed from bitset-only payloads to a generic sample-filter payload, so we should still validate no-filter and bitset search performance against `main`.

Proposed planned benchmark coverage:

- no-filter CAGRA search vs `main`
- bitset-filtered CAGRA search vs `main`
- equivalent bitset predicate vs UDF predicate
- first-call UDF compile/link latency
- warm-cache repeated UDF search latency
- representative CAGRA algorithms/configs, especially SINGLE_CTA, MULTI_CTA, and MULTI_KERNEL where applicable

This should give us both existing-functionality regression coverage and a baseline for the new UDF path.

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

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

URL: #2132

Author: dantegd

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

@@ -14,6 +14,7 @@ struct tag_i8 {};
 struct tag_u8 {};
 struct tag_filter_none {};
 struct tag_filter_bitset {};
+struct tag_filter_udf {};
 
 struct tag_bitset_u32 {};
 

cpp/include/cuvs/neighbors/cagra.hpp

@@ -346,7 +346,10 @@ struct search_params : cuvs::neighbors::search_params {
   /**
    * A parameter indicating the rate of nodes to be filtered-out, when filtering is used.
    * The value must be equal to or greater than 0.0 and less than 1.0. Default value is
-   * negative, in which case the filtering rate is automatically calculated.
+   * negative, in which case the filtering rate is automatically calculated when possible.
+   * For `filtering::udf_filter`, CAGRA uses `udf_filter::filtering_rate` when this value is
+   * negative. If both values are negative, CAGRA assumes 0.0 because a UDF's selectivity cannot be
+   * inferred from the source string.
    */
   float filtering_rate = -1.0;
 };

cpp/include/cuvs/neighbors/common.hpp

@@ -24,7 +24,9 @@
 
 #include <memory>
 #include <numeric>
+#include <string>
 #include <type_traits>
+#include <utility>
 
 #ifdef __cpp_lib_bitops
 #include <bit>
@@ -495,7 +497,7 @@ namespace filtering {
  * @{
  */
 
-enum class FilterType { None, Bitmap, Bitset };
+enum class FilterType { None, Bitmap, Bitset, UDF };
 
 struct base_filter {
   ~base_filter()                             = default;
@@ -615,6 +617,46 @@ struct bitset_filter : public base_filter {
   void to_csr(raft::resources const& handle, csr_matrix_t& csr);
 };
 
+/**
+ * @brief JIT-LTO user-defined filter predicate.
+ *
+ * The source must define a device function named by @c function_name with signature:
+ *
+ * @code{.cpp}
+ * __device__ bool cuvs_filter_udf(uint32_t query_id, source_index_t source_id, void* filter_data);
+ * @endcode
+ *
+ * Return @c true to allow a source vector to appear in the results and @c false to reject it.
+ * @c filter_data is passed through unchanged and must point to device-accessible memory when the
+ * UDF dereferences it. CAGRA currently provides @c source_index_t as @c uint32_t in the generated
+ * JIT fragment.
+ */
+struct udf_filter : public base_filter {
+  /** CUDA C++ source containing the device predicate. */
+  std::string source;
+  /** Opaque device-accessible pointer passed to the predicate. */
+  void* filter_data = nullptr;
+  /** Estimated fraction of rows rejected by the predicate, or negative if unknown. */
+  float filtering_rate = -1.0f;
+  /** Device function name to call from the generated CAGRA sample filter. */
+  std::string function_name = "cuvs_filter_udf";
+
+  udf_filter() = default;
+
+  explicit udf_filter(std::string source,
+                      void* filter_data         = nullptr,
+                      float filtering_rate      = -1.0f,
+                      std::string function_name = "cuvs_filter_udf")
+    : source(std::move(source)),
+      filter_data(filter_data),
+      filtering_rate(filtering_rate),
+      function_name(std::move(function_name))
+  {
+  }
+
+  FilterType get_filter_type() const override { return FilterType::UDF; }
+};
+
 /** @} */  // end group neighbors_filtering
 
 /**

cpp/src/neighbors/cagra.cuh

@@ -26,6 +26,8 @@
 
 #include <rmm/cuda_stream_view.hpp>
 
+#include <algorithm>
+
 namespace cuvs::neighbors::cagra {
 
 // Member function implementations for cagra::index
@@ -380,6 +382,25 @@ void search(raft::resources const& res,
     auto sample_filter_copy = sample_filter;
     return search_with_filtering<T, IdxT, decltype(sample_filter_copy), OutputIdxT>(
       res, params_copy, idx, queries, neighbors, distances, sample_filter_copy);
+  } catch (const std::bad_cast&) {
+  }
+
+  try {
+    auto& sample_filter =
+      dynamic_cast<const cuvs::neighbors::filtering::udf_filter&>(sample_filter_ref);
+    search_params params_copy = params;
+    if (params.filtering_rate < 0.0) {
+      const float min_filtering_rate = 0.0f;
+      const float max_filtering_rate = 0.999f;
+      params_copy.filtering_rate =
+        sample_filter.filtering_rate < 0.0f
+          ? 0.0f
+          : std::min(std::max(sample_filter.filtering_rate, min_filtering_rate),
+                     max_filtering_rate);
+    }
+    auto sample_filter_copy = sample_filter;
+    return search_with_filtering<T, IdxT, decltype(sample_filter_copy), OutputIdxT>(
+      res, params_copy, idx, queries, neighbors, distances, sample_filter_copy);
   } catch (const std::bad_cast&) {
     RAFT_FAIL("Unsupported sample filter type");
   }

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

@@ -0,0 +1,258 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights
+ * reserved. SPDX-License-Identifier: Apache-2.0
+ */
+#pragma once
+
+#include "../../sample_filter.cuh"  // public filter types
+#include "../sample_filter_data.cuh"
+#include "jit_lto_kernels/cagra_filter_payload.cuh"
+
+#include <raft/core/error.hpp>
+
+#include <cuda_runtime_api.h>
+
+#include <cstddef>
+#include <cstdint>
+#include <cstring>
+#include <list>
+#include <mutex>
+#include <type_traits>
+#include <unordered_map>
+
+namespace cuvs::neighbors::cagra::detail {
+
+template <typename PayloadT>
+std::uint64_t cagra_payload_hash(PayloadT const& payload)
+{
+  static_assert(std::is_trivially_copyable_v<PayloadT>);
+  constexpr std::uint64_t kOffset = 1469598103934665603ull;
+  constexpr std::uint64_t kPrime  = 1099511628211ull;
+  auto const* bytes               = reinterpret_cast<unsigned char const*>(&payload);
+  std::uint64_t hash              = kOffset;
+  for (std::size_t i = 0; i < sizeof(PayloadT); ++i) {
+    hash ^= bytes[i];
+    hash *= kPrime;
+  }
+  return hash;
+}
+
+template <typename PayloadT>
+struct cagra_device_payload_owner {
+  struct state {
+    PayloadT host_payload{};
+    PayloadT* device_payload{nullptr};
+    cudaStream_t stream{};
+    cudaEvent_t ready_event{};
+    int device{-1};
+    std::mutex mutex;
+
+    explicit state(PayloadT payload) : host_payload(payload) {}
+
+    ~state() noexcept
+    {
+      if (device_payload != nullptr) {
+        RAFT_CUDA_TRY_NO_THROW(cudaFreeAsync(device_payload, stream));
+      }
+      if (ready_event != nullptr) { RAFT_CUDA_TRY_NO_THROW(cudaEventDestroy(ready_event)); }
+    }
+
+    PayloadT* dev_ptr(cudaStream_t cuda_stream)
+    {
+      std::lock_guard<std::mutex> lock(mutex);
+      if (device_payload == nullptr) {
+        RAFT_CUDA_TRY(cudaGetDevice(&device));
+        RAFT_CUDA_TRY(cudaMallocAsync(
+          reinterpret_cast<void**>(&device_payload), sizeof(PayloadT), cuda_stream));
+        RAFT_CUDA_TRY(cudaMemcpyAsync(
+          device_payload, &host_payload, sizeof(PayloadT), cudaMemcpyHostToDevice, cuda_stream));
+        RAFT_CUDA_TRY(cudaEventCreateWithFlags(&ready_event, cudaEventDisableTiming));
+        RAFT_CUDA_TRY(cudaEventRecord(ready_event, cuda_stream));
+        stream = cuda_stream;
+      } else {
+        RAFT_CUDA_TRY(cudaStreamWaitEvent(cuda_stream, ready_event, 0));
+      }
+      return device_payload;
+    }
+  };
+
+  // PayloadT is copied to device by value. Pointer fields inside PayloadT are shallow-copied and
+  // must already point to device-addressable memory that remains valid while the cached payload is
+  // usable.
+  struct cache_key {
+    std::uint64_t payload_hash{};
+    int device{};
+
+    bool operator==(cache_key const& other) const
+    {
+      return payload_hash == other.payload_hash && device == other.device;
+    }
+  };
+
+  struct cache_key_hash {
+    std::size_t operator()(cache_key const& key) const
+    {
+      auto seed = static_cast<std::size_t>(key.payload_hash);
+      seed ^= static_cast<std::size_t>(key.device) + 0x9e3779b9 + (seed << 6) + (seed >> 2);
+      return seed;
+    }
+  };
+
+  cagra_device_payload_owner() = default;
+
+  void* dev_ptr(PayloadT payload, cudaStream_t stream) const
+  {
+    int device{};
+    RAFT_CUDA_TRY(cudaGetDevice(&device));
+
+    // Keep cached payload copies for process lifetime to avoid per-search allocation/copy churn.
+    // Cross-stream reuse is ordered by each state's ready_event before kernels consume the pointer.
+    const auto key = cache_key{cagra_payload_hash(payload), device};
+    state* selected_state{};
+    {
+      std::lock_guard<std::mutex> lock(cache_mutex_);
+      auto& entries = cache_[key];
+      for (auto& cached : entries) {
+        if (std::memcmp(&cached.host_payload, &payload, sizeof(PayloadT)) == 0) {
+          selected_state = &cached;
+          break;
+        }
+      }
+      if (selected_state == nullptr) {
+        entries.emplace_back(payload);
+        selected_state = &entries.back();
+      }
+    }
+
+    return selected_state->dev_ptr(stream);
+  }
+
+ private:
+  mutable std::mutex cache_mutex_;
+  mutable std::unordered_map<cache_key, std::list<state>, cache_key_hash> cache_;
+};
+
+template <typename T>
+struct is_bitset_filter : std::false_type {};
+
+template <typename bitset_t, typename index_t>
+struct is_bitset_filter<::cuvs::neighbors::filtering::bitset_filter<bitset_t, index_t>>
+  : std::true_type {};
+
+template <typename T>
+struct is_udf_filter : std::false_type {};
+
+template <>
+struct is_udf_filter<::cuvs::neighbors::filtering::udf_filter> : std::true_type {};
+
+template <typename SourceIndexT, typename FilterT>
+::cuvs::neighbors::detail::bitset_filter_data_t<SourceIndexT> make_cagra_bitset_filter_storage(
+  const FilterT& filter)
+{
+  const auto bitset_view = filter.view();
+  return ::cuvs::neighbors::detail::bitset_filter_data_t<SourceIndexT>{
+    const_cast<std::uint32_t*>(bitset_view.data()),
+    static_cast<SourceIndexT>(bitset_view.size()),
+    static_cast<SourceIndexT>(bitset_view.get_original_nbits())};
+}
+
+template <typename PayloadT>
+void* get_cagra_device_payload(PayloadT payload, cudaStream_t stream)
+{
+  static cagra_device_payload_owner<PayloadT> owner;
+  return owner.dev_ptr(payload, stream);
+}
+
+template <typename SourceIndexT, typename FilterT>
+void* make_cagra_bitset_filter_payload(const FilterT& filter, cudaStream_t stream)
+{
+  return get_cagra_device_payload(make_cagra_bitset_filter_storage<SourceIndexT>(filter), stream);
+}
+
+template <typename SourceIndexT, typename FilterT>
+void fill_cagra_sample_filter(cagra_sample_filter<SourceIndexT>& out,
+                              const FilterT& filter,
+                              cudaStream_t stream)
+{
+  using DecayedFilter = std::decay_t<FilterT>;
+  if constexpr (is_bitset_filter<DecayedFilter>::value) {
+    out.filter_data = make_cagra_bitset_filter_payload<SourceIndexT>(filter, stream);
+  } else if constexpr (is_udf_filter<DecayedFilter>::value) {
+    out.filter_data = filter.filter_data;
+  }
+}
+
+template <typename SourceIndexT, typename FilterT>
+std::uint64_t cagra_filter_payload_hash(const FilterT& filter)
+{
+  using DecayedFilter = std::decay_t<FilterT>;
+  if constexpr (is_bitset_filter<DecayedFilter>::value) {
+    return cagra_payload_hash(make_cagra_bitset_filter_storage<SourceIndexT>(filter));
+  } else if constexpr (requires { filter.filter; }) {
+    return cagra_filter_payload_hash<SourceIndexT>(filter.filter);
+  } else {
+    return 0;
+  }
+}
+
+template <typename FilterT>
+void* cagra_filter_data_ptr(const FilterT& filter)
+{
+  using DecayedFilter = std::decay_t<FilterT>;
+  if constexpr (is_udf_filter<DecayedFilter>::value) {
+    return filter.filter_data;
+  } else if constexpr (requires { filter.filter; }) {
+    return cagra_filter_data_ptr(filter.filter);
+  } else {
+    return nullptr;
+  }
+}
+
+template <typename SampleFilterT>
+std::uint32_t cagra_filter_query_id_offset(const SampleFilterT& sample_filter)
+{
+  if constexpr (requires {
+                  sample_filter.filter;
+                  sample_filter.offset;
+                }) {
+    return sample_filter.offset;
+  } else {
+    return 0;
+  }
+}
+
+/// Host: fill @ref cagra_sample_filter from a CAGRA filter object.
+template <typename SourceIndexT, typename SampleFilterT>
+cagra_sample_filter<SourceIndexT> extract_cagra_sample_filter(const SampleFilterT& sample_filter,
+                                                              cudaStream_t stream)
+{
+  cagra_sample_filter<SourceIndexT> out;
+  if constexpr (requires {
+                  sample_filter.filter;
+                  sample_filter.offset;
+                }) {
+    out.query_id_offset = sample_filter.offset;
+    fill_cagra_sample_filter(out, sample_filter.filter, stream);
+  } else {
+    fill_cagra_sample_filter(out, sample_filter, stream);
+  }
+  return out;
+}
+
+/// Host: find UDF compile/link metadata only. Query offsets stay in the runtime payload produced
+/// by @ref extract_cagra_sample_filter and are applied before calling the linked sample_filter.
+template <typename SampleFilterT>
+const ::cuvs::neighbors::filtering::udf_filter* get_cagra_udf_filter(
+  const SampleFilterT& sample_filter)
+{
+  using DecayedFilter = std::decay_t<SampleFilterT>;
+  if constexpr (is_udf_filter<DecayedFilter>::value) {
+    return &sample_filter;
+  } else if constexpr (requires { sample_filter.filter; }) {
+    return get_cagra_udf_filter(sample_filter.filter);
+  } else {
+    return nullptr;
+  }
+}
+
+}  // namespace cuvs::neighbors::cagra::detail