Link-time Optimization

Background

What is JIT LTO?

JIT LTO (Just-In-Time Link-Time Optimization) is a CUDA compilation strategy that enables dynamic kernel compilation and linking at runtime. Instead of pre-compiling all possible kernel variants (which would result in an explosion of binary size), JIT LTO compiles kernel fragments separately and links them together on-demand when a specific kernel configuration is needed.

Fragment Terminology

A fragment is a self-contained, compilable unit of CUDA code that can be linked with other fragments to form a complete kernel. In the JIT LTO system:

• Entrypoint Fragment: The main kernel function that serves as the entry point. This is always the __global__ kernel function.
• Device Function Fragments: Separate fragments containing device functions (e.g., distance computations, filters, post-processing) that are called by the entrypoint kernel.
• Fragment Key: A unique identifier for a fragment, typically constructed from template parameters and configuration values.
• Fatbin: The compiled binary representation of a fragment, embedded in the executable.

The key advantage is that device functions can be compiled independently and reused across multiple kernel entrypoints, reducing compilation time and binary size.

How It Works

1. Build Time: Fragments are compiled into fatbins and embedded in the executable.
2. Runtime: When a kernel needs to be launched:
   • The planner identifies which fragments are needed based on the configuration
   • Fragments are loaded from the embedded fatbins
   • Nvjitlink (Link-Time Optimization) links the fragments together
   • The linked kernel is cached and launched

Walkthrough Example

Let’s walk through creating a JIT LTO kernel system for a search kernel with templated device functions.

Step 1: Define the Kernel and Device Functions

We start with a kernel that has templated device functions that we want to separate into fragments:

search_kernel.cuh:

1
#pragma once
2
3
#include <cuda_runtime.h>
4
5
namespace example::detail {
6
7
// Device function for distance computation
8
template <typename T>
9
__device__ float compute_distance_euclidean(T a, T b) {
10
    T diff = a - b;
11
    return diff * diff;
12
}
13
14
template <typename T>
15
__device__ float compute_distance_inner_product(T a, T b) {
16
    return -a * b;  // Negative for max inner product search
17
}
18
19
// Device function for filtering
20
template <typename IdxT>
21
__device__ bool apply_filter_none(uint32_t query_id, IdxT node_id, void* filter_data) {
22
    return true;
23
}
24
25
template <typename IdxT>
26
__device__ bool apply_filter_bitset(uint32_t query_id, IdxT node_id, void* filter_data) {
27
    // Simplified - actual implementation would check bitset
28
    return true;
29
}
30
31
// Main kernel - will use generic extern device functions
32
template <typename T, typename OutT, typename IdxT, bool UseOptimizedPath, int Veclen>
33
__device__ void search_kernel_impl(
34
    const T* dataset,
35
    const T* queries,
36
    IdxT* results,
37
    OutT* distances,  // Output distance type
38
    uint32_t num_queries,
39
    uint32_t dataset_size,
40
    void* filter_data) {
41
42
    uint32_t query_id = blockIdx.x * blockDim.x + threadIdx.x;
43
    if (query_id >= num_queries) return;
44
45
    OutT best_dist = std::numeric_limits<OutT>::max();
46
    IdxT best_idx = 0;
47
48
    for (IdxT i = 0; i < dataset_size; ++i) {
49
        // Call generic extern device functions (implementations linked from fragments)
50
        if (!apply_filter<IdxT>(query_id, i, filter_data)) continue;
51
52
        OutT dist = static_cast<OutT>(compute_distance<T>(queries[query_id], dataset[i]));
53
54
        // Use optimized path if enabled
55
        if constexpr (UseOptimizedPath) {
56
            // Optimized implementation
57
            if (dist < best_dist) {
58
                best_dist = dist;
59
                best_idx = i;
60
            }
61
        } else {
62
            // Standard implementation
63
            if (dist < best_dist) {
64
                best_dist = dist;
65
                best_idx = i;
66
            }
67
        }
68
    }
69
70
    results[query_id] = best_idx;
71
    distances[query_id] = best_dist;
72
}
73
74
} // namespace example::detail

Step 2: Create Device Function Fragments

We’ll create separate header files for each device function variant. Each implements the generic function signature that the kernel expects:

compute_distance_euclidean.cuh:

1
#pragma once
2
3
namespace example::detail {
4
5
// Implements the generic compute_distance function for euclidean distance
6
template <typename T>
7
__device__ float compute_distance(T a, T b) {
8
    T diff = a - b;
9
    return diff * diff;
10
}
11
12
} // namespace example::detail

compute_distance_inner_product.cuh:

1
#pragma once
2
3
namespace example::detail {
4
5
// Implements the generic compute_distance function for inner product
6
template <typename T>
7
__device__ float compute_distance(T a, T b) {
8
    return -a * b;  // Negative for max inner product search
9
}
10
11
} // namespace example::detail

filter_none.cuh:

1
#pragma once
2
3
namespace example::detail {
4
5
// Implements the generic apply_filter function for no filtering
6
template <typename IdxT>
7
__device__ bool apply_filter(uint32_t query_id, IdxT node_id, void* filter_data) {
8
    return true;
9
}
10
11
} // namespace example::detail

filter_bitset.cuh:

1
#pragma once
2
3
namespace example::detail {
4
5
// Implements the generic apply_filter function for bitset filtering
6
template <typename IdxT>
7
__device__ bool apply_filter(uint32_t query_id, IdxT node_id, void* filter_data) {
8
    // Actual bitset implementation
9
    return true;
10
}
11
12
} // namespace example::detail

Step 3: Create JSON Matrix Files

JSON matrix files define all the parameter combinations that need to be compiled. The build system uses these to generate .cu files from .cu.in templates.

How JSON Cross-Product Works:

• The build system computes a modified Cartesian product (cross-product) of all parameter combinations.
• Leaf nodes are the actual values. These can be strings, numbers, booleans, or null, but only strings should be used, even for numbers, for example "1".
• Related values can be grouped together in a dictionary consisting of single values. Any dictionary key in such a dictionary’s ancestry will not be used in the final product, and should be prefixed with _ to indicate that it is used only for grouping.
• Keys containing only leaf nodes will be used in the final product, and should not be prefixed with _.
• The matrix product algorithm will automatically warn if the proper naming convention (_ prefix or not) is not followed.
• Each group expands to create multiple combinations, and all groups are cross-multiplied.

For example, if you have:

1
{
2
  "_data_type": [{"data_type": "float"}, {"data_type": "half"}],
3
  "_index": [{"idx_type": "uint32_t"}, {"idx_type": "int64_t"}],
4
  "capacity": ["1", "2"]
5
}

This generates 2 × 2 × 2 = 8 combinations:

• {data_type: "float", idx_type: "uint32_t", capacity: "1"}
• {data_type: "float", idx_type: "uint32_t", capacity: "2"}
• {data_type: "float", idx_type: "int64_t", capacity: "1"}
• … and so on

When a group contains nested arrays (like veclen: ["1", "4"]), those are also expanded within that group before the cross-product is computed.

compute_distance_matrix.json

1
{
2
  "_distance_type": [
3
    {
4
      "distance_name": "euclidean",
5
      "header_file": "example/jit_lto_kernels/compute_distance_euclidean.cuh"
6
    },
7
    {
8
      "distance_name": "inner_product",
9
      "header_file": "example/jit_lto_kernels/compute_distance_inner_product.cuh"
10
    }
11
  ],
12
  "_data_type": [
13
    {
14
      "data_type": "float",
15
      "type_abbrev": "f"
16
    },
17
    {
18
      "data_type": "__half",
19
      "type_abbrev": "h"
20
    }
21
  ]
22
}

filter_matrix.json

1
{
2
  "filter_name": [
3
    "filter_none",
4
    "filter_bitset"
5
  ],
6
  "_index": [
7
    {
8
      "idx_type": "uint32_t",
9
      "idx_abbrev": "ui"
10
    },
11
    {
12
      "idx_type": "int64_t",
13
      "idx_abbrev": "l"
14
    }
15
  ]
16
}

search_kernel_matrix.json

This example demonstrates conditional combinations: OutT can be float or double when T is float, but only float when T is __half.

1
{
2
  "_data_type": [
3
    {
4
      "data_type": "float",
5
      "type_abbrev": "f",
6
      "_output_type": [
7
        {
8
          "out_type": "float",
9
          "out_abbrev": "f"
10
        },
11
        {
12
          "out_type": "double",
13
          "out_abbrev": "d"
14
        }
15
      ]
16
    },
17
    {
18
      "data_type": "__half",
19
      "type_abbrev": "h",
20
      "_output_type": [
21
        {
22
          "out_type": "float",
23
          "out_abbrev": "f"
24
        }
25
      ]
26
    }
27
  ],
28
  "_index": [
29
    {
30
      "idx_type": "uint32_t",
31
      "idx_abbrev": "ui"
32
    },
33
    {
34
      "idx_type": "int64_t",
35
      "idx_abbrev": "l"
36
    }
37
  ],
38
  "_optimized": [
39
    {
40
      "optimized_name": "optimized",
41
      "optimized_value": "true",
42
      "veclen": ["1", "4"]
43
    },
44
    {
45
      "optimized_name": "standard",
46
      "optimized_value": "false",
47
      "veclen": ["8", "16"]
48
    }
49
  ]
50
}

This generates 24 combinations (3 data/output type combinations × 2 index types × 4 optimized/veclen combinations):

• float + float + uint32_t + optimized + veclen=1
• float + float + uint32_t + optimized + veclen=4
• float + float + uint32_t + standard + veclen=8
• float + float + uint32_t + standard + veclen=16
• float + double + uint32_t + optimized + veclen=1
• float + double + uint32_t + optimized + veclen=4
• float + double + uint32_t + standard + veclen=8
• float + double + uint32_t + standard + veclen=16
• __half + float + uint32_t + optimized + veclen=1
• __half + float + uint32_t + optimized + veclen=4
• __half + float + uint32_t + standard + veclen=8
• __half + float + uint32_t + standard + veclen=16
• … and the same with int64_t (total: 24 combinations)

Step 4: Create .cu.in Template Files

The .cu.in files are templates that get instantiated for each combination in the JSON matrix. They contain explicit template instantiations.

compute_distance_kernel.cu.in

/*
 * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
#include "@header_file@"
namespace example::detail {
// Instantiate the generic compute_distance device function template
// The specific implementation (euclidean or inner_product) comes from the header
template __device__ float compute_distance<@data_type@>(@data_type@, @data_type@);
} // namespace example::detail

filter_kernel.cu.in

/*
 * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
#include "example/jit_lto_kernels/@filter_name@.cuh"
namespace example::detail {
// Instantiate the generic apply_filter device function template
// The specific implementation (filter_none or filter_bitset) comes from the header
template __device__ bool apply_filter<@idx_type@>(uint32_t, @idx_type@, void*);
} // namespace example::detail

Update search_kernel.cuh with Extern Declarations

The kernel header needs to declare generic extern device functions so the kernel code can call them. The specific implementations will be linked from fragments at runtime:

search_kernel.cuh:

1
#pragma once
2
3
#include <cuda_runtime.h>
4
5
namespace example::detail {
6
7
// Forward declare generic extern device functions that will be linked from fragments
8
// The specific implementations (euclidean, inner_product, etc.) are resolved at link time
9
template <typename T>
10
extern __device__ float compute_distance(T, T);
11
12
template <typename IdxT>
13
extern __device__ bool apply_filter(uint32_t, IdxT, void*);
14
15
// Main kernel - uses generic extern device functions
16
template <typename T, typename OutT, typename IdxT, bool UseOptimizedPath, int Veclen>
17
__device__ void search_kernel_impl(
18
    const T* dataset,
19
    const T* queries,
20
    IdxT* results,
21
    OutT* distances,  // Output distance type
22
    uint32_t num_queries,
23
    uint32_t dataset_size,
24
    void* filter_data) {
25
26
    uint32_t query_id = blockIdx.x * blockDim.x + threadIdx.x;
27
    if (query_id >= num_queries) return;
28
29
    OutT best_dist = std::numeric_limits<OutT>::max();
30
    IdxT best_idx = 0;
31
32
    for (IdxT i = 0; i < dataset_size; ++i) {
33
        // Call generic extern device functions (specific implementations linked from fragments)
34
        if (!apply_filter<IdxT>(query_id, i, filter_data)) continue;
35
36
        OutT dist = static_cast<OutT>(compute_distance<T>(queries[query_id], dataset[i]));
37
38
        // Use optimized path if enabled
39
        if constexpr (UseOptimizedPath) {
40
            // Optimized implementation
41
            if (dist < best_dist) {
42
                best_dist = dist;
43
                best_idx = i;
44
            }
45
        } else {
46
            // Standard implementation
47
            if (dist < best_dist) {
48
                best_dist = dist;
49
                best_idx = i;
50
            }
51
        }
52
    }
53
54
    results[query_id] = best_idx;
55
    distances[query_id] = best_dist;
56
}
57
58
} // namespace example::detail

search_kernel.cu.in

The .cu.in file only contains the explicit template instantiation:

/*
 * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
#include "example/jit_lto_kernels/search_kernel.cuh"
namespace example::detail {
// Instantiate the kernel template
extern "C" __global__ void search_kernel(
    const @data_type@* dataset, const @data_type@* queries, @idx_type@* results, @out_type@* distances,
    uint32_t num_queries, uint32_t dataset_size, void* filter_data)
{
  search_kernel_impl<@data_type@, @out_type@, @idx_type@, @optimized_value@, @veclen@>(
    dataset,
    queries,
    results,
    distances,
    num_queries,
    dataset_size,
    filter_data);
}
} // namespace example::detail

Note: The kernel uses generic function templates (compute_distance<T> and apply_filter<IdxT>) that are resolved at link time. The specific implementations (euclidean vs inner_product, filter_none vs filter_bitset) are provided by the fragments that get linked together.

Step 5: Create Fragment Tags for Embedding

Fragment tags register the compiled fatbins so they can be loaded at runtime. They are used to help the linker find and include the relevant fatbins at build time. When calling generate_jit_lto_kernels(), we pass a FRAGMENT_TAG_FORMAT argument, which constructs the tag type from the given placeholders, and a FRAGMENT_TAG_HEADER_FILES argument, which specifies one or more header files that the fragment tags come from. The JIT+LTO system will then automatically generate and compile a .cpp file that registers the fragment using the provided tag.

Important: When requesting fragments from the AlgorithmPlanner, we use tags (like tag_f, tag_h) instead of real types (like float, __half) in the add_static_fragment template parameters. This avoids including heavy headers that define the actual types, significantly improving compilation times. The tags are lightweight empty structs that serve only as compile-time identifiers.

registration_tags.hpp

1
#pragma once
2
3
struct tag_h{};
4
struct tag_f{};
5
struct tag_d{};
6
struct tag_ui{};
7
struct tag_l{};
8
9
struct tag_metric_euclidean {};
10
struct tag_metric_inner_product {};
11
12
struct tag_filter_none {};
13
struct tag_filter_bitset {};
14
15
template <typename DataTag, typename OutTag, typename IdxTag, bool Optimized, int Veclen>
16
struct fragment_tag_search {};
17
18
template <typename DistanceTag, typename DataTag>
19
struct fragment_tag_compute_distance {};
20
21
template <typename FilterTag, typename IndexTag>
22
struct fragment_tag_filter {};

Step 6: Create the Planner

The planner is responsible for:

1. Identifying which fragments are needed for a given configuration
2. Building a unique key for the fragment combination
3. Requesting the fragments from the fragment database
4. Linking them together to create a launchable kernel

CRITICAL: The fragment keys constructed in the planner methods must match EXACTLY with the keys used in the corresponding FRAGMENT_TAG_FORMAT argument. Any mismatch will result in runtime linking failures.

search_planner.hpp:

1
#pragma once
2
3
#include <cuvs/detail/jit_lto/AlgorithmPlanner.hpp>
4
#include <cuvs/detail/jit_lto/MakeFragmentKey.hpp>
5
#include <cuvs/detail/jit_lto/registration_tags.hpp>
6
#include <string>
7
8
struct SearchPlanner : AlgorithmPlanner {
9
  inline static LauncherJitCache launcher_jit_cache{};
10
11
  SearchPlanner()
12
    : AlgorithmPlanner("search_kernel", launcher_jit_cache)
13
  {
14
  }
15
16
  template <typename DataTag, typename OutTag, typename IdxTag, bool Optimized, int Veclen>
17
  void add_search_function()
18
  {
19
    add_static_fragment<fragment_tag_search<DataTag, OutTag, IdxTag, Optimized, Veclen>>();
20
  }
21
22
  template <typename DistanceTag, typename DataTag>
23
  void add_compute_distance_function()
24
  {
25
    add_static_fragment<fragment_tag_compute_distance<DistanceTag, DataTag>>();
26
  }
27
28
  template <typename FilterTag, typename IndexTag>
29
  void add_filter_function()
30
  {
31
    add_static_fragment<fragment_tag_filter<FilterTag, IndexTag>>();
32
  }
33
};

Step 7: Integrate with Code Path

Now we integrate the planner into the actual search function:

search_jit.cuh:

1
#pragma once
2
3
#include "search_planner.hpp"
4
#include <cuvs/detail/jit_lto/registration_tags.hpp>
5
#include <raft/core/device_resources.hpp>
6
7
namespace example::detail {
8
9
// Type tag helpers
10
template <typename T>
11
constexpr auto get_data_type_tag() {
12
  if constexpr (std::is_same_v<T, float>) return tag_f{};
13
  if constexpr (std::is_same_v<T, __half>) return tag_h{};
14
}
15
16
template <typename IdxT>
17
constexpr auto get_idx_type_tag() {
18
  if constexpr (std::is_same_v<IdxT, uint32_t>) return tag_ui{};
19
  if constexpr (std::is_same_v<IdxT, int64_t>) return tag_l{};
20
}
21
22
template <typename OutType>
23
constexpr auto get_out_type_tag() {
24
  if constexpr (std::is_same_v<OutType, float>) return tag_f{};
25
  if constexpr (std::is_same_v<OutType, double>) return tag_d{};
26
}
27
28
template <DistanceType Metric>
29
constexpr auto get_metric_tag() {
30
  if constexpr (Metric == DistanceType::Euclidean) return tag_metric_euclidean{};
31
  if constexpr (Metric == DistanceType::InnerProduct) return tag_metric_inner_product{};
32
}
33
34
template <FilterType Filter>
35
constexpr auto get_filter_tag() {
36
  if constexpr (Filter == FilterType::None) return tag_filter_none{};
37
  if constexpr (Filter == FilterType::Bitset) return tag_filter_bitset{};
38
}
39
40
template <typename T, typename OutT, typename IdxT, DistanceType Metric, FilterType Filter, bool Optimized, int Veclen>
41
void search_jit(
42
    raft::device_resources const& handle,
43
    const T* dataset,
44
    const T* queries,
45
    IdxT* results,
46
    OutT* distances,
47
    uint32_t num_queries,
48
    uint32_t dataset_size,
49
    void* filter_data = nullptr) {
50
51
  using data_tag = decltype(get_data_type_tag<T>());
52
  using idx_tag = decltype(get_idx_type_tag<IdxT>());
53
  using out_tag = decltype(get_out_type_tag<OutT>());
54
  using metric_tag = decltype(get_metric_tag<Metric>());
55
  using filter_tag = decltype(get_filter_tag<Filter>());
56
57
  // Create planner with type tags and boolean parameter
58
  // Note: The boolean is appended to the fragment key since make_fragment_key
59
  // cannot handle non-type template parameters
60
  SearchPlanner planner;
61
62
  // Add required device function fragments
63
  planner.add_search_function<data_tag, out_tag, idx_tag, Optimized, Veclen>();
64
  planner.add_compute_distance_device_function<metric_tag, data_tag>();
65
  planner.add_filter_device_function<filter_tag, idx_tag>();
66
67
  // Get the launcher (this will build/link fragments if needed)
68
  auto launcher = planner.get_launcher();
69
70
  // Launch configuration
71
  dim3 block(256);
72
  dim3 grid((num_queries + block.x - 1) / block.x);
73
74
  // Launch the kernel - arguments are passed directly
75
  launcher->dispatch(
76
      raft::resource::get_cuda_stream(handle),
77
      grid,
78
      block,
79
      0,  // shared memory size
80
      dataset,
81
      queries,
82
      results,
83
      distances,
84
      num_queries,
85
      dataset_size,
86
      filter_data);
87
}
88
89
} // namespace example::detail

Key Concepts

Fragment Tags

Fragment tags uniquely identify fragments. They’re simple lightweight types that are passed as the sole template parameter to StaticFatbinFragmentEntry:

1
template <typename OutT>
2
struct fragment_tag_get_score {};

Fragment tags may themselves take template parameters in order to uniquely identify them. Typically, one fragment tag template will correspond to a single function, and a fragment tag template specialization will correspond to a function specialization.

When a fatbin is compiled and embedded in C++ code, a translation unit specializes StaticFatbinFragmentEntry to specify its data and length static fields:

1
using _FragmentEntry = StaticFatbinFragmentEntry<fragment_tag_get_score<uint32_t>>;
2
3
template <>
4
const uint8_t* const _FragmentEntry::data = embedded_fatbin;
5
6
template <>
7
const size_t _FragmentEntry::length = sizeof(embedded_fatbin);

Then, an AlgorithmPlanner can call add_static_fragment() with the fragment tag (NOT the StaticFatbinFragmentEntry specialization) as the sole template parameter:

1
template <typename OutTag>
2
void add_get_score_function()
3
{
4
  add_static_fragment<fragment_tag_get_score<OutTag>>();
5
}

At build time, the linker takes care of finding and including the static fragments that have been specified by the algorithm planner.

Registration Tags

Registration tags are type-safe identifiers used to organize fragments. They’re typically empty structs:

1
struct tag_f {};  // float
2
struct tag_h {};  // half
3
struct tag_ui {}; // uint32_t
4
struct tag_l {};  // int64_t

These tags are used in registerAlgorithm<>() to create a hierarchical organization of fragments.

AlgorithmLauncher

The AlgorithmLauncher is the runtime handle for a linked kernel. It:

• Holds a cudaKernel_t handle to the linked kernel
• Provides call() and call_cooperative() methods to launch the kernel
• Manages the lifetime of the cudaLibrary_t that contains the kernel

Best Practices

1. Minimize Includes: JIT LTO fragments should have minimal includes, especially avoiding host-side headers. Extract device-only code into separate headers.

2. Fragment Granularity: Balance between too many small fragments (overhead) and too few large fragments (less reuse). Device functions that are reused across multiple kernels are good candidates for separate fragments.

3. Naming Consistency: Ensure fragment tags match exactly between registration and lookup. Use helper functions to construct tags consistently.

4. Type Safety: Use registration tags to provide compile-time type safety and avoid runtime string mismatches.

5. Caching: Each planner type should hold a static LauncherJitCache and pass it to AlgorithmPlanner; get_launcher() then reuses linked kernels for the same fragment key within that cache.

Example: IVF Flat

IVF Flat uses JIT LTO with:

• Metric fragments: Euclidean and inner product distance computations (16 fatbins)
• Post-lambda fragments: Identity, sqrt, and compose post-processing (3 fatbins)
• Interleaved scan fragments: Main search kernel with various configurations (320 fatbins)
• Filter fragments: None and bitset filters (2 fatbins)

Total: 341 fatbins that can be combined into many more kernel variants at runtime.

Step 8: Integrate with CMake Build System

To integrate JIT LTO kernels into the CMake build system, add calls to generate_jit_lto_kernels() in your main CMakeLists.txt file (typically in cpp/CMakeLists.txt).

The generate_jit_lto_kernels() function (defined in cmake/modules/generate_jit_lto_kernels.cmake) takes:

• NAME_FORMAT: Format string for generated kernel names (using @variable@ syntax)
• MATRIX_JSON_FILE: Path to the JSON matrix file
• KERNEL_INPUT_FILE: Path to the .cu.in template
• FRAGMENT_TAG_FORMAT: Format string for fragment tag type (using @variable@ syntax)
• FRAGMENT_TAG_HEADER_FILES: List of header files that provide the fragment tag types (can be enclosed in </> or "/", automatically enclosed in quotes if quotes and brackets are not provided)
• OUTPUT_DIRECTORY: Where generated files are placed
• KERNEL_LINK_LIBRARIES: Interface library with compilation settings

Call generate_jit_lto_kernels() once for each fragment type (compute_distance, filter, search_kernel, etc.). The function reads the JSON matrix, computes the cross-product of all combinations, generates .cu and .cpp files from the templates, compiles them into fatbins, and returns a list of generated source files that should be added to your JIT LTO library target.

See the CUVS cpp/CMakeLists.txt file for a complete example of how to set up the interface library, call generate_jit_lto_kernels() for each fragment type, and create the final library target.

Summary

JIT LTO enables:

• Reduced binary size: Compile fragments once, combine many ways
• Faster compilation: Fragments compile independently
• Runtime flexibility: Link fragments on-demand based on configuration
• Code reuse: Device function fragments shared across kernels

The process involves:

1. Separating device functions into fragment headers
2. Creating JSON matrices defining parameter combinations
3. Creating .cu.in templates for explicit instantiations
4. Creating fragment tag types for fatbin registration
5. Creating a planner to manage fragment dependencies
6. Integrating the planner into the code path to launch kernels
7. Adding CMake integration to generate and compile all fragment variants

Fragment Architecture

JIT LTO kernels are split into fragments, which are fatbins containing individual pieces of code that can be strung together rather than instantiating the whole kernel at once. Each fragment only needs to be multiplied out over the dimensions (template parameters) that the fragment itself contains rather than the kernel as a whole. At runtime, these fragments are combined together by nvjitlink into the final program.

In JIT LTO, there are two kinds of code: algorithms and adapters. Algorithms are, roughly speaking, code that actually “does stuff” - searching, sorting, even as simple as initializing variables. Adapters don’t do anything by themselves, but are merely thin wrappers around algorithms that exist only for reducing the number of template parameters that the caller needs to know about. It should generally be assumed that algorithm code is expensive to multiply over a matrix, and thus such multiplication should be minimized, while adapter code is cheap to multiply.

An algorithm function is a function that contains real code for the algorithm, and an adapter function merely calls an algorithm function with more template parameters than the adapter function itself has. An algorithm file contains algorithm code, and an adapter file contains adapter code.

Here is an example of an algorithm file that contains an algorithm function:

1
template <typename T, T Divisor>
2
__device__ bool is_divisible_impl(T value)
3
{
4
  return value % Divisor == 0;
5
}

Here is an example of an adapter file that contains an adapter function:

#include "device_functions.cuh"  // is_divisible
#include "is_divisible_impl.cuh" // is_divisible_impl
namespace {
using data_t             = @data_type@;
constexpr data_t divisor = @divisor@;
}  // namespace
template <>
__device__ bool is_divisible<data_t>(data_t value)
{
  return is_divisible_impl<data_t, divisor>(value);
}

This is the most common pattern that you will see in cuVS’s JIT LTO code. Note that any code that calls is_divisible() does not need to know the value of Divisor, which allows the caller to be multiplied over fewer dimensions, thus reducing the amount of code generated.

Note that in the above adapter file, @data_type@ and @divisor@ are build-time substitutions performed by CMake. These substitutions will be filled in with values from the matrix product. Note that they are all grouped together in a single namespace, making it easy to find all substitutions. This should be preferred to sprinkling the substitutions throughout the code.

Here is an example with two algorithm files:

1
// greater_than_impl.cuh
2
#include "device_impl_functions.cuh" // filter
3
4
template <typename T, T Comparand>
5
__device__ bool filter(T value)
6
{
7
  return value > Comparand;
8
}

1
// less_than_impl.cuh
2
#include "device_impl_functions.cuh" // filter
3
4
template <typename T, T Comparand>
5
__device__ bool filter(T value)
6
{
7
  return value < Comparand;
8
}

And here is the accompanying adapter file:

#include "@op_name@_impl.cuh" // filter
namespace {
using data_t = @data_type@;
}
template __device__ bool filter<data_t>(data_t value);

This is another common pattern that you will see in cuVS JIT LTO. Note that the adapter file does not contain any adapter functions, but merely instantiates a different algorithm function based on which algorithm file is included based on the CMake substitution.

When a piece of algorithm code is used in multiple kernels, it should be split into its own shared fragment. At this point, it becomes important to also distinguish algorithm fragments and adapter fragments. An algorithm fragment contains an algorithm function that exposes all of the relevant template parameters, and this fragment is shared between multiple kernels. An adapter fragment is specific to a kernel. If a kernel wishes to invoke the same shared algorithm multiple times in the same run with different template parameters, it can employ multiple adapter fragments to accomplish this. Consider the following header file:

1
// filter.cuh
2
3
template <typename T, T Comparand>
4
__device__ bool filter_less_than(T value);
5
6
template <typename T, T Comparand>
7
__device__ bool filter_greater_than(T value);

And the following adapter files:

#include "device_functions.cuh" // filter_first_pass
#include "filter.cuh"           // filter
namespace {
using data_t               = @data_type@;
constexpr data_t comparand = @comparand@;
}
template <>
__device__ bool filter_first_pass<data_t>(data_t value)
{
  return filter_@op_name@<data_t, comparand>(value);
}

#include "device_functions.cuh" // filter_second_pass
#include "filter.cuh"           // filter
namespace {
using data_t               = @data_type@;
constexpr data_t comparand = @comparand@;
}
template <>
__device__ bool filter_second_pass<data_t>(data_t value)
{
  return filter_@op_name@<data_t, comparand>(value);
}

And the following algorithm file:

1
#include "device_functions.cuh" // filter_first_pass, filter_second_pass
2
3
template <typename T>
4
__device__ bool filter_all_passes(T value)
5
{
6
  return filter_first_pass<T>(value) && filter_second_pass<T>(value);
7
}

Note that filter_first_pass and filter_second_pass both invoke one of the filter functions, but which one they invoke is decided independently for each. Also note that neither of the adapter fragments contains the underlying algorithm code, but rather links against the corresponding shared algorithm fragments.

The key to minimizing code generation is to minimize the number of dimensions that any given fragment needs to be multiplied out over. If a section of algorithm code uses lots of template parameters, try to separate out sections that use only a subset of these parameters, put them into their own fragment, and remove the corresponding template parameters from the caller. Make judicious use of adapter code to accomplish this. An adapter function should only have the template parameters that appear in its signature, whereas an algorithm function should have all of the template parameters that appear in its signature or its implementation.

Unoptimized algorithm:

1
#include "filter_less_than.cuh"
2
3
template <typename T, T Comparand>
4
__device__ size_t find_first(T* values, size_t count)
5
{
6
  for (size_t i = 0; i < count; i++) {
7
    if (filter_less_than_impl<T, Comparand>(values[i])) {
8
      return i;
9
    }
10
  }
11
12
  // Could not find any
13
  return count;
14
}

Note that the algorithm includes the Comparand template parameter, which means the entire algorithm has to be multiplied out over all the possible values of this parameter.

Optimized algorithm:

1
#include "device_functions.cuh"
2
3
template <typename T>
4
__device__ size_t find_first(T* values, size_t count)
5
{
6
  for (size_t i = 0; i < count; i++) {
7
    if (filter_less_than<T>(values[i])) {
8
      return i;
9
    }
10
  }
11
12
  // Could not find any
13
  return count;
14
}

We are now using an adapter function (possibly inside an adapter fragment) called filter_less_than to invoke filter_less_than_impl (which may be inside a shared algorithm fragment). This allows us to hide the Comparand parameter from find_first, which means we no longer need to multiply the entire algorithm over all possible values of Comparand, only the filter_less_than adapter and algorithm.