UDF Usage

Caution: Custom distance metrics for IVF-flat search are experimental. They live under the cuvs::neighbors::ivf_flat::experimental::udf namespace and the associated CUVS_METRIC macro. APIs and behavior may change without a major release.

What this feature does

You can supply your own CUDA device code that defines how distance accumulates between a query vector and database vectors inside the IVF-flat interleaved scan (the fine search over lists). Technical background on compilation and linking is in Link-time Optimization.

Available via C++ APIs

• IVF-flat — search (search_params.metric_udf / CUVS_METRIC).

Requirements and tips

• Include <cuvs/neighbors/ivf_flat.hpp> and define a metric with CUVS_METRIC(MyName, { ... }). Set search_params.metric_udf to the string returned by MyName_udf().
• Prefer the helpers when combining lanes so one body works for scalar and packed int8_t / uint8_t as well as wider element types.
• Custom UDF is not supported for fp16 (__half / half) indices at this time; the headers enforce this with a static assertion when applicable.
• The scan assumes ascending distance order for top-k selection; metrics that do not behave like a distance in that sense need careful validation.
• The first search with a new metric string may pay a one-time compilation cost; reuse the same string (and run a warmup) to benefit from the caches described in JIT Compilation.

Example

1
#include <cuvs/neighbors/ivf_flat.hpp>
2
3
namespace ivf = cuvs::neighbors::ivf_flat;
4
5
// L∞ (Chebyshev): per dimension, acc = max(acc, |x - y|); acc starts at 0 in the scan kernel.
6
CUVS_METRIC(my_chebyshev, {
7
  auto d = abs_diff(x, y);
8
  acc    = (d > acc) ? d : acc;
9
})
10
11
void run_search(raft::resources const& res,
12
                ivf::index<float, int64_t> const& index,
13
                raft::device_matrix_view<const float, int64_t, raft::row_major> queries,
14
                raft::device_matrix_view<int64_t, int64_t, raft::row_major> neighbors,
15
                raft::device_matrix_view<float, int64_t, raft::row_major> distances)
16
{
17
  ivf::search_params params;
18
  params.metric_udf = my_chebyshev_udf();
19
20
  ivf::search(res, params, index, queries, neighbors, distances);
21
}

Helpers in CUVS_METRIC bodies

Inside CUVS_METRIC(MyName, { ... }) you write the body of operator()(AccT& acc, point_type x, point_type y). In scope: acc, x, y, template parameters T, AccT, Veclen, and the helpers below. The macro’s full argument list and notes live beside CUVS_METRIC in <cuvs/neighbors/ivf_flat.hpp>.

More examples: cpp/tests/neighbors/ann_ivf_flat/test_udf.cu.

Further reading

• C++ API reference: neighbors::ivf_flat
• JIT LTO architecture and IVF-flat fragments: Link-time Optimization