From 5b13cb755287fc609d5d94756be6842052c8c604 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Thu, 16 Apr 2026 02:15:47 +0000 Subject: [PATCH 1/5] udf docs --- docs/source/advanced_topics.rst | 7 + docs/source/jit_lto_guide.md | 221 +++++++++++++++++++++++++++++++- docs/source/udf_usage.rst | 71 ++++++++++ 3 files changed, 298 insertions(+), 1 deletion(-) create mode 100644 docs/source/udf_usage.rst diff --git a/docs/source/advanced_topics.rst b/docs/source/advanced_topics.rst index 4171845af5..cac539c144 100644 --- a/docs/source/advanced_topics.rst +++ b/docs/source/advanced_topics.rst @@ -2,6 +2,7 @@ Advanced Topics =============== - `Just-in-Time Compilation`_ +- :doc:`UDF Usage ` Just-in-Time Compilation ------------------------ @@ -16,7 +17,13 @@ Thus, the JIT compilation is a one-time cost and you can expect no loss in real Currently, the following capabilities will trigger a JIT compilation: - IVF Flat search APIs: :doc:`cuvs::neighbors::ivf_flat::search() ` +UDFs are available in the following APIs: +----------------------------------------- +- IVF Flat search (C++ only): experimental custom distance via ``search_params.metric_udf``; see + :doc:`udf_usage`. + .. toctree:: :maxdepth: 2 jit_lto_guide + udf_usage diff --git a/docs/source/jit_lto_guide.md b/docs/source/jit_lto_guide.md index 490bec5914..9250910560 100644 --- a/docs/source/jit_lto_guide.md +++ b/docs/source/jit_lto_guide.md @@ -573,8 +573,10 @@ The planner is responsible for: #pragma once #include +#include #include #include +#include #include struct SearchPlanner : AlgorithmPlanner { @@ -602,6 +604,16 @@ struct SearchPlanner : AlgorithmPlanner { { add_static_fragment>(); } + + void add_metric_udf_fragment(std::unique_ptr fragment) + { + add_fragment(std::move(fragment)); + } + + void add_filter_udf_fragment(std::unique_ptr fragment) + { + add_fragment(std::move(fragment)); + } }; ``` @@ -617,9 +629,13 @@ Now we integrate the planner into the actual search function: #include "search_planner.hpp" #include #include +#include namespace example::detail { +enum class DistanceType { Euclidean }; +enum class FilterType { None }; + // Type tag helpers template constexpr auto get_data_type_tag() { @@ -671,7 +687,6 @@ void search_jit( // cannot handle non-type template parameters SearchPlanner planner; - // Add required device function fragments planner.add_search_function(); planner.add_compute_distance_device_function(); planner.add_filter_device_function(); @@ -701,6 +716,210 @@ void search_jit( } // namespace example::detail ``` +### Step 7b: Example — NVRTC UDFs for `compute_distance` and `apply_filter` + +Same entry kernel as Steps 1–7, but `compute_distance` / `apply_filter` are **not** linked from static matrix fatbins: one NVRTC TU per hook (or compile twice) and register each through the planner’s **UDF-specific** APIs (see below—not the same calls as static matrix fragments). Both are **templates** in this example, so each TU must include a **forwarding definition** of the hook plus an **explicit instantiation** for every concrete specialization the entry fatbin calls (e.g. `compute_distance` and `apply_filter`). + +**1. Entry / shared header — declarations only** + +```cpp +namespace example::detail { + +template +__device__ float compute_distance(T q, T d); + +template +__device__ bool apply_filter(uint32_t query_id, IdxT node_id, void* filter_data); + +} // namespace example::detail +``` + +Register each NVRTC fatbin through your planner’s **UDF** hooks (the snippet below uses `add_metric_udf_fragment` / `add_filter_udf_fragment`), not through the static-matrix helpers from Step 6 (`add_compute_distance_function` / `add_filter_function` on `SearchPlanner`). Those paths select different embedded fatbins; using both for the same hook would duplicate device definitions. + +If you only declare `template ... compute_distance` / `apply_filter` in the entry TU, the NVRTC program must define them **and** emit matching `template __device__ ...` **explicit instantiations** for the concrete types the entry calls. For `IdxT`-dependent hooks, do not bake `IdxT` into a macro string: append those lines from a small host helper (`instantiate_apply_filter_udf` below) that takes the same NVRTC spelling `type_name()` returns, analogous to `instantiate_compute_distance_udf` for `T`. + +**2. Write the body first, then a `#define` turns it into NVRTC source** + +The hooks are **function-like macros** (`#define EXAMPLE_UDF_DISTANCE(NAME, BODY) ...`). You only edit the braced **body** argument. The preprocessor emits (1) a real `__device__` function template so NVCC checks `T`, `q`, `d`, etc., and (2) a `NAME_udf()` / `NAME_filter_udf()` factory that concatenates boilerplate with `std::string(#BODY)` so NVRTC compiles the **same** tokens NVCC already parsed. If the body needs raw `"` characters, splice that part with raw-string concatenation instead of relying on `#BODY` alone. For distance here, the NVRTC text also defines `compute_distance_udf_impl`, which forwards to `NAME_distance`. + +**Macro definitions** (typically a shared header included before the invocations): + +```cpp +#include +#include + +#define EXAMPLE_PP_CAT_(a, b) a##b +#define EXAMPLE_PP_CAT(a, b) EXAMPLE_PP_CAT_(a, b) +#define EXAMPLE_PP_STR_(x) #x +#define EXAMPLE_PP_STR(x) EXAMPLE_PP_STR_(x) + +// NAME_udf(): NVRTC program defines NAME_distance and compute_distance_udf_impl; host appends +// instantiate_compute_distance_udf (forwarding compute_distance + explicit inst only). +#define EXAMPLE_UDF_DISTANCE(NAME, BODY) \ + template \ + __device__ float EXAMPLE_PP_CAT(NAME, _distance)(T q, T d) BODY \ + \ + inline std::string EXAMPLE_PP_CAT(NAME, _udf)() \ + { \ + return std::string("#include \n" \ + "namespace example::detail {\n" \ + "template \n" \ + "__device__ float " EXAMPLE_PP_STR(EXAMPLE_PP_CAT(NAME, _distance)) \ + "(T q, T d) ") \ + + std::string(#BODY) + \ + std::string("\n" \ + "template \n" \ + "__device__ float compute_distance_udf_impl(T q, T d) {\n" \ + " return ") + \ + std::string(EXAMPLE_PP_STR(EXAMPLE_PP_CAT(NAME, _distance))) + \ + std::string("(q, d);\n" \ + "}\n}\n"); \ + } + +// Forwarding compute_distance + explicit inst only (user metric stays inside NAME_udf()). +inline std::string instantiate_compute_distance_udf(char const* t_type) +{ + std::ostringstream oss; + oss << "\nnamespace example::detail {\n" + << "template \n" + << "__device__ float compute_distance(T q, T d) {\n" + << " return compute_distance_udf_impl(q, d);\n" + << "}\n" + << "template __device__ float compute_distance<" << t_type << ">(" << t_type << ", " << t_type + << ");\n" + << "}\n"; + return oss.str(); +} + +// Device NAME_filter + filter_udf(); append instantiate_apply_filter_udf for apply_filter. +#define EXAMPLE_UDF_FILTER(NAME, BODY) \ + template \ + __device__ bool EXAMPLE_PP_CAT(NAME, _filter)(uint32_t query_id, IdxT node_id, void* filter_data) \ + BODY \ + \ + inline std::string EXAMPLE_PP_CAT(NAME, _filter_udf)() \ + { \ + return std::string("#include \n" \ + "namespace example::detail {\n" \ + "template \n" \ + "__device__ bool " EXAMPLE_PP_STR(EXAMPLE_PP_CAT(NAME, _filter)) \ + "(uint32_t query_id, IdxT node_id, void* filter_data) ") \ + + std::string(#BODY) + \ + std::string("\n" \ + "template \n" \ + "__device__ bool apply_filter(uint32_t query_id, IdxT node_id, " \ + "void* filter_data) {\n" \ + " return " EXAMPLE_PP_STR(EXAMPLE_PP_CAT(NAME, _filter)) \ + "(query_id, node_id, filter_data);\n" \ + "}\n" \ + "}\n"); \ + } + +// Call after NAME_filter_udf() string is concatenated, before compile. +inline std::string instantiate_apply_filter_udf(char const* idx_type) +{ + std::ostringstream oss; + oss << "\nnamespace example::detail {\n" + << "template __device__ bool apply_filter<" << idx_type << ">(uint32_t, " << idx_type + << ", void*);\n" + << "}\n"; + return oss.str(); +} +``` + +**Invocations** (same CUDA source or header — each line is a **function-like macro call** that expands into a device template plus a `std::string` factory; no extra `#define` is required unless you want a named alias): + +```cpp +EXAMPLE_UDF_DISTANCE(my_l2, { + T diff = q - d; + return diff * diff; +}) + +EXAMPLE_UDF_FILTER(my_pass, { + (void)query_id; + (void)node_id; + (void)filter_data; + return true; +}) +``` + +`#BODY` follows the usual preprocessor rules (avoid raw `"` inside the body unless you splice that section with raw-string literals around `#BODY`). + +**3. Host — one NVRTC compile per UDF fragment** + +Assemble the full CUDA program for each hook, compile it once, and register the fatbin through the planner’s UDF entry point (e.g. `add_metric_udf_fragment`). Do not merge unrelated UDFs into a single compile / fragment. + +Step 7b’s toy kernel uses `compute_distance` and `apply_filter` (explicit lines come from `instantiate_compute_distance_udf` / `instantiate_apply_filter_udf`, not from hard-coded types inside the macros). A matching `type_name` only has to spell those concrete types; each return must be the **exact** token NVRTC will parse (same characters as in the generated CUDA). Strip cv/ref first, then add `if constexpr` branches as you support more index and element types. + +```cpp +#include + +template +constexpr const char* type_name() +{ + using T = std::remove_cv_t>; + if constexpr (std::is_same_v) { + return "float"; + } else if constexpr (std::is_same_v) { + return "uint32_t"; + } else { + static_assert(std::is_same_v, "add a branch for each T / AccT / IdxT the entry uses"); + return ""; + } +} +``` + +Call `instantiate_compute_distance_udf` / `instantiate_apply_filter_udf` to append forwarding templates and explicit instantiations, using the same `type_name<...>()` tokens the entry TU will use. Concatenate that glue onto each `*_udf()` string, then compile and register. + +**4. Extend `search_jit.cuh` (Step 7) for UDF vs static** + +Step 7 only registered static matrix fragments. To add NVRTC UDFs without breaking the Euclidean / no-filter path, add `#include `, replace the single-value `DistanceType` / `FilterType` enums and the `get_metric_tag` / `get_filter_tag` templates with the extended versions below, add the empty UDF tag structs, and keep the UDF glue (`my_l2_udf`, `instantiate_*`, `type_name`, `nvrtc_compiler`) in the same translation unit. Then **replace** the two unconditional `add_compute_distance_device_function` / `add_filter_device_function` lines with the `if constexpr` planner block (second snippet). + +```cpp +enum class DistanceType { Euclidean, MetricUdf }; +enum class FilterType { None, FilterUdf }; + +struct tag_metric_custom_udf {}; +struct tag_filter_custom_udf {}; + +template +constexpr auto get_metric_tag() { + if constexpr (Metric == DistanceType::Euclidean) return tag_metric_euclidean{}; + else if constexpr (Metric == DistanceType::MetricUdf) return tag_metric_custom_udf{}; +} + +template +constexpr auto get_filter_tag() { + if constexpr (Filter == FilterType::None) return tag_filter_none{}; + else if constexpr (Filter == FilterType::FilterUdf) return tag_filter_custom_udf{}; +} +``` + +```cpp +SearchPlanner planner; +planner.add_search_function(); + +if constexpr (std::is_same_v) { + std::string metric_udf_code = my_l2_udf(); + metric_udf_code += instantiate_compute_distance_udf(type_name()); + planner.add_metric_udf_fragment(nvrtc_compiler().compile(metric_udf_code, metric_udf_code)); +} else { + planner.add_compute_distance_device_function(); +} + +if constexpr (std::is_same_v) { + std::string filter_udf_code = my_pass_filter_udf(); + filter_udf_code += instantiate_apply_filter_udf(type_name()); + planner.add_filter_udf_fragment(nvrtc_compiler().compile(filter_udf_code, filter_udf_code)); +} else { + planner.add_filter_device_function(); +} + +auto launcher = planner.get_launcher(); +``` + +Instantiate `search_jit` with `DistanceType::MetricUdf` and/or `FilterType::FilterUdf` only when you intend the NVRTC branches; `Euclidean` and `FilterType::None` keep the original static behavior. + ## Key Concepts ### Fragment Tags diff --git a/docs/source/udf_usage.rst b/docs/source/udf_usage.rst new file mode 100644 index 0000000000..b984f0a2a7 --- /dev/null +++ b/docs/source/udf_usage.rst @@ -0,0 +1,71 @@ +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 :doc:`jit_lto_guide`. + +Available via C++ APIs for the following algorithms +--------------------------------------------------- + +* IVF-flat — :doc:`search ` (``search_params.metric_udf`` / ``CUVS_METRIC``). + +Requirements and tips +----------------------- + +* Include ```` and define a metric with ``CUVS_METRIC(MyName, { ... })``. + Set ``search_params.metric_udf`` to the string returned by ``MyName_udf()``. +* Prefer the helpers documented next to the macro (``squared_diff``, ``abs_diff``, ``dot_product``, + ``point`` element access, and so on) so the same definition works across ``float``, ``int8_t`` / + ``uint8_t`` packed lanes, and related accumulator 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 :doc:`advanced_topics`. + +Example +------- + +.. code-block:: cpp + + #include + + namespace ivf = cuvs::neighbors::ivf_flat; + + // L∞ (Chebyshev): per dimension, acc = max(acc, |x - y|); acc starts at 0 in the scan kernel. + CUVS_METRIC(my_chebyshev, { + auto d = abs_diff(x, y); + acc = (d > acc) ? d : acc; + }) + + void run_search(raft::resources const& res, + ivf::index const& index, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances) + { + ivf::search_params params; + params.metric_udf = my_chebyshev_udf(); + + ivf::search(res, params, index, queries, neighbors, distances); + } + +For more examples (L2 via ``squared_diff``, raw string fragments, and so on), see +``cpp/tests/neighbors/ann_ivf_flat/test_udf.cu`` in the cuVS repository. + +Further reading +--------------- + +* C++ API reference: :doc:`cpp_api/neighbors_ivf_flat` +* JIT LTO architecture and IVF-flat fragments: :doc:`jit_lto_guide` From adb7dcb3f67a8b725c6d00302926806ea6a1583e Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Tue, 12 May 2026 20:49:56 +0000 Subject: [PATCH 2/5] address reviews --- cpp/include/cuvs/neighbors/ivf_flat.hpp | 6 +- docs/source/jit_lto_guide.md | 97 ++++++++++++++++++------- docs/source/udf_usage.rst | 39 ++++++++-- 3 files changed, 106 insertions(+), 36 deletions(-) diff --git a/cpp/include/cuvs/neighbors/ivf_flat.hpp b/cpp/include/cuvs/neighbors/ivf_flat.hpp index f214db295c..2255d84865 100644 --- a/cpp/include/cuvs/neighbors/ivf_flat.hpp +++ b/cpp/include/cuvs/neighbors/ivf_flat.hpp @@ -3523,10 +3523,8 @@ inline std::string instantiate_udf(char const* data_type, char const* acc_type, * }) * * CUVS_METRIC(my_chebyshev, { - * for (int i = 0; i < x.size(); ++i) { - * auto diff = (x[i] > y[i]) ? (x[i] - y[i]) : (y[i] - x[i]); - * if (diff > acc) acc = diff; - * } + * auto d = abs_diff(x, y); + * acc = (d > acc) ? d : acc; * }) */ #define CUVS_METRIC(NAME, BODY) \ diff --git a/docs/source/jit_lto_guide.md b/docs/source/jit_lto_guide.md index 9250910560..134878db9b 100644 --- a/docs/source/jit_lto_guide.md +++ b/docs/source/jit_lto_guide.md @@ -605,6 +605,7 @@ struct SearchPlanner : AlgorithmPlanner { add_static_fragment>(); } + // Same as add_fragment(std::move(fragment)); distinct names are for readability at call sites. void add_metric_udf_fragment(std::unique_ptr fragment) { add_fragment(std::move(fragment)); @@ -657,12 +658,20 @@ constexpr auto get_out_type_tag() { template constexpr auto get_metric_tag() { - if constexpr (Metric == DistanceType::Euclidean) return tag_metric_euclidean{}; + if constexpr (Metric == DistanceType::Euclidean) { + return tag_metric_euclidean{}; + } else { + static_assert(!sizeof(Metric*), "extend get_metric_tag when adding DistanceType enumerators"); + } } template constexpr auto get_filter_tag() { - if constexpr (Filter == FilterType::None) return tag_filter_none{}; + if constexpr (Filter == FilterType::None) { + return tag_filter_none{}; + } else { + static_assert(!sizeof(Filter*), "extend get_filter_tag when adding FilterType enumerators"); + } } template @@ -718,9 +727,32 @@ void search_jit( ### Step 7b: Example — NVRTC UDFs for `compute_distance` and `apply_filter` -Same entry kernel as Steps 1–7, but `compute_distance` / `apply_filter` are **not** linked from static matrix fatbins: one NVRTC TU per hook (or compile twice) and register each through the planner’s **UDF-specific** APIs (see below—not the same calls as static matrix fragments). Both are **templates** in this example, so each TU must include a **forwarding definition** of the hook plus an **explicit instantiation** for every concrete specialization the entry fatbin calls (e.g. `compute_distance` and `apply_filter`). +**What you’re building.** The same search kernel as Steps 1–7 still calls `compute_distance` / `apply_filter`, but for a UDF build those symbols are **not** taken from prebuilt matrix fatbins: you compile a small NVRTC program per hook at runtime and register it with the planner so LTO links it next to the entry fragment. + +**How the pieces connect.** + +```mermaid +flowchart LR + subgraph entry["Entry fatbin"] + K["Kernel calls templates"] + H["Header: declare only"] + end + subgraph nvrtc["Per-hook NVRTC TU"] + M["Macro: device body + string factory"] + G["Host glue: forwarding + explicit inst"] + end + subgraph plan["Planner"] + R["add_*_udf_fragment(fatbin)"] + end + K --> H + H --> M + M --> G + G --> R +``` + +**1. Shared header — forward declarations** -**1. Entry / shared header — declarations only** +The entry TU matches Step 1–7: templates are declared here and defined elsewhere at link time. ```cpp namespace example::detail { @@ -734,15 +766,11 @@ __device__ bool apply_filter(uint32_t query_id, IdxT node_id, void* filter_data) } // namespace example::detail ``` -Register each NVRTC fatbin through your planner’s **UDF** hooks (the snippet below uses `add_metric_udf_fragment` / `add_filter_udf_fragment`), not through the static-matrix helpers from Step 6 (`add_compute_distance_function` / `add_filter_function` on `SearchPlanner`). Those paths select different embedded fatbins; using both for the same hook would duplicate device definitions. - -If you only declare `template ... compute_distance` / `apply_filter` in the entry TU, the NVRTC program must define them **and** emit matching `template __device__ ...` **explicit instantiations** for the concrete types the entry calls. For `IdxT`-dependent hooks, do not bake `IdxT` into a macro string: append those lines from a small host helper (`instantiate_apply_filter_udf` below) that takes the same NVRTC spelling `type_name()` returns, analogous to `instantiate_compute_distance_udf` for `T`. - -**2. Write the body first, then a `#define` turns it into NVRTC source** +**2. NVRTC source — macros and string factories** -The hooks are **function-like macros** (`#define EXAMPLE_UDF_DISTANCE(NAME, BODY) ...`). You only edit the braced **body** argument. The preprocessor emits (1) a real `__device__` function template so NVCC checks `T`, `q`, `d`, etc., and (2) a `NAME_udf()` / `NAME_filter_udf()` factory that concatenates boilerplate with `std::string(#BODY)` so NVRTC compiles the **same** tokens NVCC already parsed. If the body needs raw `"` characters, splice that part with raw-string concatenation instead of relying on `#BODY` alone. For distance here, the NVRTC text also defines `compute_distance_udf_impl`, which forwards to `NAME_distance`. +Use **function-like macros** so you edit only the `{ ... }` body; the preprocessor still emits a real `__device__` template for NVCC, and `NAME_udf()` / `NAME_filter_udf()` build the CUDA text NVRTC compiles. The distance macro also emits `compute_distance_udf_impl` calling `NAME_distance`; host-side `instantiate_compute_distance_udf` only appends the forwarding `compute_distance` plus its explicit instantiation (same idea as `instantiate_apply_filter_udf` for `apply_filter`). -**Macro definitions** (typically a shared header included before the invocations): +**Macro definitions** (shared header; include before the invocations): ```cpp #include @@ -827,7 +855,7 @@ inline std::string instantiate_apply_filter_udf(char const* idx_type) } ``` -**Invocations** (same CUDA source or header — each line is a **function-like macro call** that expands into a device template plus a `std::string` factory; no extra `#define` is required unless you want a named alias): +**Invocations** (file scope — each call expands the macro once): ```cpp EXAMPLE_UDF_DISTANCE(my_l2, { @@ -843,13 +871,11 @@ EXAMPLE_UDF_FILTER(my_pass, { }) ``` -`#BODY` follows the usual preprocessor rules (avoid raw `"` inside the body unless you splice that section with raw-string literals around `#BODY`). +Avoid raw `"` inside `#BODY` unless you splice that part with raw-string concatenation around `#BODY`. -**3. Host — one NVRTC compile per UDF fragment** +**3. Host — `type_name`, glue, compile, register** -Assemble the full CUDA program for each hook, compile it once, and register the fatbin through the planner’s UDF entry point (e.g. `add_metric_udf_fragment`). Do not merge unrelated UDFs into a single compile / fragment. - -Step 7b’s toy kernel uses `compute_distance` and `apply_filter` (explicit lines come from `instantiate_compute_distance_udf` / `instantiate_apply_filter_udf`, not from hard-coded types inside the macros). A matching `type_name` only has to spell those concrete types; each return must be the **exact** token NVRTC will parse (same characters as in the generated CUDA). Strip cv/ref first, then add `if constexpr` branches as you support more index and element types. +Each full NVRTC program is one string: `*_udf()` plus `instantiate_*` output. `type_name()` must return the **exact** token the entry TU uses (e.g. `float`, `uint32_t`). ```cpp #include @@ -863,19 +889,19 @@ constexpr const char* type_name() } else if constexpr (std::is_same_v) { return "uint32_t"; } else { - static_assert(std::is_same_v, "add a branch for each T / AccT / IdxT the entry uses"); + static_assert(std::is_same_v, "add a branch for each concrete T / IdxT you use"); return ""; } } ``` -Call `instantiate_compute_distance_udf` / `instantiate_apply_filter_udf` to append forwarding templates and explicit instantiations, using the same `type_name<...>()` tokens the entry TU will use. Concatenate that glue onto each `*_udf()` string, then compile and register. - -**4. Extend `search_jit.cuh` (Step 7) for UDF vs static** +**4. Planner — extend Step 7 for UDF vs static** -Step 7 only registered static matrix fragments. To add NVRTC UDFs without breaking the Euclidean / no-filter path, add `#include `, replace the single-value `DistanceType` / `FilterType` enums and the `get_metric_tag` / `get_filter_tag` templates with the extended versions below, add the empty UDF tag structs, and keep the UDF glue (`my_l2_udf`, `instantiate_*`, `type_name`, `nvrtc_compiler`) in the same translation unit. Then **replace** the two unconditional `add_compute_distance_device_function` / `add_filter_device_function` lines with the `if constexpr` planner block (second snippet). +Step 7 used only static fragments. Add `#include `, extend enums/tags/`get_*_tag` as below, keep UDF glue in the same TU as `search_jit`, then swap the two unconditional `add_compute_distance_device_function` / `add_filter_device_function` calls for this block: ```cpp +// Widen config: static Euclidean vs NVRTC metric, static none vs NVRTC filter. +// Extend the DistanceType / FilterType enums from Step 7: enum class DistanceType { Euclidean, MetricUdf }; enum class FilterType { None, FilterUdf }; @@ -884,14 +910,24 @@ struct tag_filter_custom_udf {}; template constexpr auto get_metric_tag() { - if constexpr (Metric == DistanceType::Euclidean) return tag_metric_euclidean{}; - else if constexpr (Metric == DistanceType::MetricUdf) return tag_metric_custom_udf{}; + if constexpr (Metric == DistanceType::Euclidean) { + return tag_metric_euclidean{}; + } else if constexpr (Metric == DistanceType::MetricUdf) { + return tag_metric_custom_udf{}; + } else { + static_assert(!sizeof(Metric*), "extend get_metric_tag when adding DistanceType enumerators"); + } } template constexpr auto get_filter_tag() { - if constexpr (Filter == FilterType::None) return tag_filter_none{}; - else if constexpr (Filter == FilterType::FilterUdf) return tag_filter_custom_udf{}; + if constexpr (Filter == FilterType::None) { + return tag_filter_none{}; + } else if constexpr (Filter == FilterType::FilterUdf) { + return tag_filter_custom_udf{}; + } else { + static_assert(!sizeof(Filter*), "extend get_filter_tag when adding FilterType enumerators"); + } } ``` @@ -899,6 +935,7 @@ constexpr auto get_filter_tag() { SearchPlanner planner; planner.add_search_function(); +// Metric: NVRTC TU vs prebuilt matrix fragment (Step 6 helpers). if constexpr (std::is_same_v) { std::string metric_udf_code = my_l2_udf(); metric_udf_code += instantiate_compute_distance_udf(type_name()); @@ -918,7 +955,13 @@ if constexpr (std::is_same_v) { auto launcher = planner.get_launcher(); ``` -Instantiate `search_jit` with `DistanceType::MetricUdf` and/or `FilterType::FilterUdf` only when you intend the NVRTC branches; `Euclidean` and `FilterType::None` keep the original static behavior. +Use `DistanceType::MetricUdf` / `FilterType::FilterUdf` only when you want the NVRTC branches; otherwise keep `Euclidean` / `None` for the original static path. + +> **Pitfalls and constraints** +> +> * **Do not** register the same hook through both UDF APIs (`add_metric_udf_fragment` / `add_filter_udf_fragment`) and the Step 6 static helpers (`add_compute_distance_function` / `add_filter_function`): they pull different fatbins and you will duplicate device definitions. +> * The NVRTC program must **define** every template the entry calls **and** emit matching **`template __device__ ...` explicit instantiations** for each concrete specialization (e.g. `compute_distance`, `apply_filter`). Prefer small host helpers (`instantiate_*` + `type_name`) for type spellings instead of hard-coding index types inside macro strings. +> * One NVRTC compile per logical TU; do not concatenate unrelated UDFs into one program string. ## Key Concepts diff --git a/docs/source/udf_usage.rst b/docs/source/udf_usage.rst index b984f0a2a7..d34abc2538 100644 --- a/docs/source/udf_usage.rst +++ b/docs/source/udf_usage.rst @@ -24,9 +24,8 @@ Requirements and tips * Include ```` and define a metric with ``CUVS_METRIC(MyName, { ... })``. Set ``search_params.metric_udf`` to the string returned by ``MyName_udf()``. -* Prefer the helpers documented next to the macro (``squared_diff``, ``abs_diff``, ``dot_product``, - ``point`` element access, and so on) so the same definition works across ``float``, ``int8_t`` / - ``uint8_t`` packed lanes, and related accumulator types. +* Prefer :ref:`udf-metric-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 @@ -61,8 +60,38 @@ Example ivf::search(res, params, index, queries, neighbors, distances); } -For more examples (L2 via ``squared_diff``, raw string fragments, and so on), see -``cpp/tests/neighbors/ann_ivf_flat/test_udf.cu`` in the cuVS repository. +.. _udf-metric-helpers: + +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 +````. + +.. list-table:: + :header-rows: 1 + :widths: 28 72 + + * - Helper + - Role + * - ``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 --------------- From 920c00e3efa3096b75b5f8321b750a6e07fdf8eb Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Tue, 12 May 2026 22:42:57 +0000 Subject: [PATCH 3/5] text instead of mermaid --- docs/source/jit_lto_guide.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/source/jit_lto_guide.md b/docs/source/jit_lto_guide.md index eb11651b5d..79be37b33f 100644 --- a/docs/source/jit_lto_guide.md +++ b/docs/source/jit_lto_guide.md @@ -676,9 +676,9 @@ void search_jit( **What you’re building.** The same search kernel as Steps 1–7 still calls `compute_distance` / `apply_filter`, but for a UDF build those symbols are **not** taken from prebuilt matrix fatbins: you compile a small NVRTC program per hook at runtime and register it with the planner so LTO links it next to the entry fragment. -**How the pieces connect.** +**How the pieces connect** (arrows read left to right): -```mermaid +```text flowchart LR subgraph entry["Entry fatbin"] K["Kernel calls templates"] From feb5ab6a8acca62c9949640bd9c578cfbd0aa98c Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Thu, 21 May 2026 01:19:14 +0000 Subject: [PATCH 4/5] cleanup --- fern/pages/advanced_topics.md | 6 ------ 1 file changed, 6 deletions(-) diff --git a/fern/pages/advanced_topics.md b/fern/pages/advanced_topics.md index 98df488a4e..6b5682a26b 100644 --- a/fern/pages/advanced_topics.md +++ b/fern/pages/advanced_topics.md @@ -12,9 +12,3 @@ Use these pages when working on specialized cuVS development topics that need lo - [JIT Compilation](jit_compilation.md): understand when cuVS triggers just-in-time compilation, how caches behave, and how to warm up JIT-compiled kernels. - [Link-time Optimization](jit_lto_guide.md): use JIT LTO for CUDA compilation, fragment generation, and runtime linking workflows. - [UDF Usage](udf_usage.md): supply custom CUDA distance metrics for IVF-flat search (C++ only, experimental). - -## User-defined metrics (UDFs) - -UDFs are available in the following APIs: - -- IVF-flat search (C++ only): experimental custom distance via `search_params.metric_udf`; see [UDF Usage](udf_usage.md). From 7c4317f68a0c634dd93acd4e8c489c5b82ea943c Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Thu, 21 May 2026 18:44:57 +0000 Subject: [PATCH 5/5] more moving around --- fern/docs.yml | 22 ++++++++++---------- fern/pages/advanced_topics.md | 6 ++---- fern/pages/cpp_api/cpp-api-cluster-kmeans.md | 16 ++++++++++---- fern/pages/developer_guide.md | 10 ++++----- fern/pages/udf_usage.md | 6 +----- fern/pages/user_guide.md | 9 ++++++++ 6 files changed, 40 insertions(+), 29 deletions(-) diff --git a/fern/docs.yml b/fern/docs.yml index 47e5d863f2..7e42d246ab 100644 --- a/fern/docs.yml +++ b/fern/docs.yml @@ -194,6 +194,13 @@ navigation: path: "./pages/user_guide/abi_stability.md" - page: "Integration Patterns" path: "./pages/user_guide/integration_patterns.md" + - section: "Advanced Topics" + path: "./pages/advanced_topics.md" + contents: + - page: "JIT Compilation" + path: "./pages/jit_compilation.md" + - page: "UDF Usage" + path: "./pages/udf_usage.md" - page: "References" path: "./pages/references.md" - section: "Developer Guide" @@ -210,17 +217,10 @@ navigation: path: "./pages/java_guidelines.md" - page: "Python Guidelines" path: "./pages/python_guidelines.md" - - section: "Advanced Topics" - path: "./pages/advanced_topics.md" - contents: - - page: "ABI Stability" - path: "./developer_guide/abi_stability.md" - - page: "JIT Compilation" - path: "./pages/jit_compilation.md" - - page: "Link-time Optimization" - path: "./pages/jit_lto_guide.md" - - page: "UDF Usage" - path: "./pages/udf_usage.md" + - page: "ABI Stability" + path: "./developer_guide/abi_stability.md" + - page: "Link-time Optimization" + path: "./pages/jit_lto_guide.md" - page: "Contributing" path: "./pages/contributing.md" - section: "API Reference" diff --git a/fern/pages/advanced_topics.md b/fern/pages/advanced_topics.md index 6b5682a26b..41eb8d8f6d 100644 --- a/fern/pages/advanced_topics.md +++ b/fern/pages/advanced_topics.md @@ -1,14 +1,12 @@ --- -slug: developer-guide/advanced-topics +slug: user-guide/advanced-topics --- # Advanced Topics -Use these pages when working on specialized cuVS development topics that need lower-level implementation context, runtime behavior, or platform-specific guidance. +Use these pages when you need lower-level implementation context, runtime behavior, or specialized cuVS features beyond the standard API surface. ## Topic Guides -- [ABI Stability](../developer_guide/abi_stability.md): understand ABI expectations for developer-facing APIs and downstream integrations. - [JIT Compilation](jit_compilation.md): understand when cuVS triggers just-in-time compilation, how caches behave, and how to warm up JIT-compiled kernels. -- [Link-time Optimization](jit_lto_guide.md): use JIT LTO for CUDA compilation, fragment generation, and runtime linking workflows. - [UDF Usage](udf_usage.md): supply custom CUDA distance metrics for IVF-flat search (C++ only, experimental). diff --git a/fern/pages/cpp_api/cpp-api-cluster-kmeans.md b/fern/pages/cpp_api/cpp-api-cluster-kmeans.md index 0979c1bb0f..fe12a7d565 100644 --- a/fern/pages/cpp_api/cpp-api-cluster-kmeans.md +++ b/fern/pages/cpp_api/cpp-api-cluster-kmeans.md @@ -48,8 +48,8 @@ struct params : base_params { ... }; | `oversampling_factor` | `double` | Oversampling factor for use in the k-means\|\| algorithm | | `batch_samples` | `int` | batch_samples and batch_centroids are used to tile 1NN computation which is useful to optimize/control the memory footprint Default tile is [batch_samples x n_clusters] i.e. when batch_centroids is 0 then don't tile the centroids NB: These parameters are unrelated to streaming_batch_size, which controls how many samples to transfer from host to device per batch when processing out-of-core data. | | `batch_centroids` | `int` | if 0 then batch_centroids = n_clusters | -| `init_size` | `int64_t` | Number of samples to randomly draw for the KMeansPlusPlus initialization step. A random subset of this size is used for centroid seeding. Only applies when dataset is on host; for device data the full dataset is always used for seeding and this parameter is ignored. When set to 0 (default) with host data uses `min(3 * n_clusters, n_samples)` as a default. Default: 0. | -| `streaming_batch_size` | `int64_t` | Number of samples to process per GPU batch when fitting with host data. When set to 0, defaults to n_samples (process all at once). Only used by the batched (host-data) code path and ignored by device-data overloads. Default: 0 (process all data at once). | +| `init_size` | `int64_t` | Number of samples to randomly draw for the KMeansPlusPlus initialization step. A random subset of this size is used for centroid seeding. Only applies when dataset is on host; for device data the full dataset is always used for seeding and this parameter is ignored. When set to 0 (default) with host data uses `min(3 * n_clusters, n_samples)` as a default. In Batched multi-GPU host-data fits, the effective KMeansPlusPlus initialization sample is materialized on device on every rank. Every rank must have enough GPU memory for this sample, and rank 0 must also have enough GPU memory for the seeding workspace. Default: 0. | +| `streaming_batch_size` | `int64_t` | Number of samples to process per GPU batch when fitting with host data. When set to 0, defaults to n_samples (process all at once). Only used by the batched (host-data) code path and ignored by device-data overloads. In multi-GPU mode, this is a per-rank batch size. Each rank processes up to this many local samples per batch, clamped to that rank's local sample count. Default: 0 (process all data at once). | ### cluster::kmeans::balanced_params @@ -108,13 +108,21 @@ raft::host_scalar_view n_iter); TODO: Evaluate replacing the extent type with int64_t. Reference issue: https://github.com/rapidsai/cuvs/issues/1961 -This overload supports out-of-core computation where the dataset resides on the host. Data is processed in GPU-sized batches, streaming from host to device. The batch size is controlled by params.streaming_batch_size. +This overload supports out-of-core computation where the dataset resides on the host. Data is processed in GPU-sized batches, streaming from host to device. The batch size is controlled by params.streaming_batch_size. In multi-GPU mode, this is a per-rank batch size. + +Multi-GPU dispatch is selected automatically based on the handle state: + +- If `raft::resource::is_multi_gpu(handle)` (cuVS SNMG): the full dataset X is split across GPUs internally with an OpenMP parallel region and NCCL. +- If `raft::resource::comms_initialized(handle)` (Dask/Ray/MPI): X is treated as this worker's partition, and RAFT communicators are used for collectives. +- Otherwise: single-GPU batched k-means. + +With `params.init == InitMethod::KMeansPlusPlus` in multi-GPU mode, the effective initialization sample must fit in GPU memory on every rank because it is materialized on every device. Rank 0 must also have enough GPU memory for the seeding workspace before centroids are broadcast. **Parameters** | Name | Direction | Type | Description | | --- | --- | --- | --- | -| `handle` | in | `raft::resources const&` | The raft handle. | +| `handle` | in | `raft::resources const&` | The raft handle. When a multi-GPU resource is attached, multi-GPU dispatch is used automatically. | | `params` | in | [`const cuvs::cluster::kmeans::params&`](/api-reference/cpp-api-cluster-kmeans#cluster-kmeans-params) | Parameters for KMeans model. Batch size is read from params.streaming_batch_size. | | `X` | in | `raft::host_matrix_view` | Training instances on HOST memory. The data must be in row-major format. [dim = n_samples x n_features] | | `sample_weight` | in | `std::optional>` | Optional weights for each observation in X (on host). [len = n_samples] | diff --git a/fern/pages/developer_guide.md b/fern/pages/developer_guide.md index e3a9a018dd..657baeacb4 100644 --- a/fern/pages/developer_guide.md +++ b/fern/pages/developer_guide.md @@ -14,13 +14,13 @@ Use these pages when contributing to cuVS or working on integrations that need t - [Java Guidelines](java_guidelines.md): follow cuVS Java API design, native binding, packaging, and resource-management conventions. - [Python Guidelines](python_guidelines.md): follow cuVS Python API, packaging, and binding conventions. -## Advanced Topics +## ABI Stability -- [Advanced Topics](advanced_topics.md): find specialized development topics and low-level implementation guidance. - [ABI Stability](../developer_guide/abi_stability.md): understand ABI expectations for developer-facing APIs and downstream integrations. -- [JIT Compilation](jit_compilation.md): understand when cuVS triggers just-in-time compilation and how runtime caches behave. -- [Link-time Optimization](jit_lto_guide.md): use JIT LTO for CUDA compilation and linking workflows. -- [UDF Usage](udf_usage.md): supply custom CUDA distance metrics for IVF-flat search (C++ only, experimental). + +## Link-time Optimization + +- [Link-time Optimization](jit_lto_guide.md): use JIT LTO for CUDA compilation, fragment generation, and runtime linking workflows. ## Contributing diff --git a/fern/pages/udf_usage.md b/fern/pages/udf_usage.md index d660819211..137add871c 100644 --- a/fern/pages/udf_usage.md +++ b/fern/pages/udf_usage.md @@ -1,7 +1,3 @@ ---- -slug: developer-guide/advanced-topics/udf-usage ---- - # 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. @@ -48,7 +44,7 @@ void run_search(raft::resources const& res, } ``` -## Helpers in `CUVS_METRIC` bodies {#helpers-in-cuvs_metric-bodies} +## 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 ``. diff --git a/fern/pages/user_guide.md b/fern/pages/user_guide.md index 7a8be2e2e5..50933ca3fd 100644 --- a/fern/pages/user_guide.md +++ b/fern/pages/user_guide.md @@ -53,4 +53,13 @@ Use these guides when you are ready to apply cuVS APIs, benchmark algorithms, or - [Compatibility](user_guide/abi_stability.md): understand cuVS release compatibility, ABI windows, and stable binary boundaries. - [Integration Patterns](user_guide/integration_patterns.md): compare direct, offloaded, and service-oriented ways to integrate cuVS into products. + +## Advanced Topics + +- [Advanced Topics](advanced_topics.md): find specialized usage topics and low-level implementation guidance. +- [JIT Compilation](jit_compilation.md): understand when cuVS triggers just-in-time compilation and how runtime caches behave. +- [UDF Usage](udf_usage.md): supply custom CUDA distance metrics for IVF-flat search (C++ only, experimental). + +## References + - [References](references.md): cite the research papers behind cuVS vector search, preprocessing, clustering, and GPU primitives.