User GuideField Guide

UDF Usage

View as Markdown

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
3namespace ivf = cuvs::neighbors::ivf_flat;
4
5// L∞ (Chebyshev): per dimension, acc = max(acc, |x - y|); acc starts at 0 in the scan kernel.
6CUVS_METRIC(my_chebyshev, {
7  auto d = abs_diff(x, y);
8  acc    = (d > acc) ? d : acc;
9})
10
11void 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>.

HelperRole
point (x, y)Element view: raw(), operator[](i), size(), is_packed().
squared_diff(x, y)Squared difference; typical building block for L2-style energy.
abs_diff(x, y)Absolute difference per lane.
dot_product(x, y)Dot product / packed-byte dot where applicable.
product(x, y)Element-wise product.
sum(x, y)Element-wise sum.
max_elem(x, y)Element-wise maximum.

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

Further reading