CUB Tuning Infrastructure

Device-scope algorithms in CUB have many knobs that do not affect the algorithms’ correctness but can significantly impact performance. For instance, the number of threads per block and items per thread can be tuned to maximize performance for a given device and data type. This document describes CUB’s tuning Infrastructure, a set of tools facilitating the process of selecting optimal tuning parameters for a given device and data type.

Definitions

We omit the word “tuning” but assume it in the definitions for all terms below, so those terms may mean something else in a more generic context.

Algorithms are tuned for different workloads, which are sub spaces of all benchmark versions defined by NVBench via a benchmark’s axes. For instance, radix sort can be tuned for different key types, different number of keys, and different distributions of keys. We partition the space spanned by a benchmark’s axes into two categories:

Compile-time (ct) Workload - a workload that can be recognized at compile time. For instance, the combination of key type and offset type is a compile-time workload for radix sort. A compile-time workload is a point in the space spanned by the Cartesian product of compile-time type axes of NVBench.
Runtime (rt) Workload - a workload that can be recognized only at runtime. For instance, the number of keys along with their distribution is a runtime workload for radix sort. A runtime workload is a point in the space spanned by the Cartesian product of non-compile-time type axes of NVBench.

The tuning infrastructure can optimize algorithms only for specific compile-time workloads, aggregating results across all runtime workloads: It searches through a space of parameters to find the combination for a given compile-time workload with the highest score:

Parameter - a parameter that can be tuned to maximize performance for a given device and data type. For instance, the number of threads per block and items per thread are tuning parameters.
Parameter Space - the set of all possible values for a given tuning parameter. Parameter Space is specific to algorithm. For instance, the parameter space for the number of threads per block is \(\{32, 64, 96, 128, \dots, 1024\}\) for radix sort, but \(\{32, 64, 128, 256, 512\}\) for merge sort.
Parameter Point - a concrete value of a tuning parameter. For instance, the parameter point for the number of threads per block is \(threads\_per\_block=128\).
Search Space - Cartesian product of parameter spaces. For instance, search space for an algorithm with tunable items per thread and threads per block might look like \(\{(ipt \times tpb) | ipt \in \{1, \dots, 25\} \text{and} tpb \in \{32, 64, 96, 128, \dots, 1024\}\}\).
Variant - a point in the corresponding search space.
Base - the variant that CUB uses by default.
Score - a single number representing the performance for a given compile-time workload across all runtime workloads. For instance, a weighted-sum of speedups of a given variant compared to its base for all runtime workloads is a score.
Search - a process consisting of covering all variants for all compile-time workloads to find a variant with maximal score.

Authoring Benchmarks

CUB benchmarks are split into multiple files based on the algorithm they are testing and potentially further into compile-time flavors that are tuned for individually (e.g.: sorting only keys vs. key-value pairs, or reducing using sum vs. using min). The name of the directory represents the name of the algorithm. The filename corresponds on the flavor. For instance, the benchmark benchmarks/bench/radix_sort/keys.cu tests the radix sort implementation sorting only keys.0 The file name is going to be transformed into cub.bench.radix_sort.keys..., which is the benchmark name reported by the infrastructure.

Benchmarks are based on NVBench. You start writing a benchmark by including nvbench_helper.cuh. It contains all necessary includes and definitions.

#include <nvbench_helper.cuh>

The next step is to define a search space. The search space is represented by a number of C++ comments. The format consists of the %RANGE% keyword, a parameter macro, a short parameter name, and a range. The range is represented by three numbers: start:end:step. Start and end are included. For instance, the following code defines a search space for two parameters, the number of threads per block and items per thread.

// %RANGE% TUNE_ITEMS_PER_THREAD ipt 7:24:1
// %RANGE% TUNE_THREADS_PER_BLOCK tpb 128:1024:32

Next, you need to define a benchmark function. The function accepts nvbench::state &state and a nvbench::type_list. For more details on the benchmark signature, take a look at the NVBench documentation.

template <typename T, typename OffsetT>
void algname(nvbench::state &state, nvbench::type_list<T, OffsetT>)
{

Tuning relies on CUB’s device algorithms to expose a dispatch layer which can be parameterized by a policy hub. CUB usually provides a default policy hub, but when tuning we want to overwrite it, so we have to specialize the dispatch layer. The tuning infrastructure will use the TUNE_BASE macro to distinguish between compiling the base version (i.e. baseline) of a benchmark and compiling a variant for a given set of tuning parameters. When base is used, no policy is specified, so that the default one CUB provides is used. If TUNE_BASE is not defined, we specify a custom policy using the parameter macros defined in the %RANGE% comments which specify the search space.

#if TUNE_BASE
  using dispatch_t = cub::DispatchReduce<T, OffsetT>; // uses default policy hub
#else
  template <typename AccumT, typename OffsetT>
  struct policy_hub_t {
    struct MaxPolicy : cub::ChainedPolicy<300, policy_t, policy_t> {
      static constexpr int threads_per_block  = TUNE_THREADS_PER_BLOCK;
      static constexpr int items_per_thread   = TUNE_ITEMS_PER_THREAD;
      ...
    };
  };

  using dispatch_t = cub::DispatchReduce<T, OffsetT, policy_hub_t<accum_t, offset_t>>;
#endif

The custom policy hub used for tuning should only expose a single MaxPolicy so CUB will use it. It must contain all parameters from the search space.

The state passed into the benchmark function allows access to runtime workload axes, for example the number of elements to process. When creating containers for the input avoid to initialize data yourself. Instead, use the gen function, which will fill the input vector with random data on GPU with no compile-time overhead.

const auto elements = static_cast<std::size_t>(state.get_int64("Elements{io}"));
thrust::device_vector<T> in(elements);
thrust::device_vector<T> out(1);

gen(seed_t{}, in);

In addition to benchmark runtime, NVBench can also report information on the achieved memory bandwidth. For this, you can optionally provide information on the memory reads and writes of the algorithm to the state:

state.add_element_count(elements);
state.add_global_memory_reads<T>(elements, "Size");
state.add_global_memory_writes<T>(1);

Most CUB algorithms need to be called twice:

once to query the amount of temporary storage needed,
once to run the actual algorithm.

We perform the first call now and allocate temporary storage:

std::size_t temp_size;
dispatch_t::Dispatch(nullptr,
                     temp_size,
                     d_in,
                     d_out,
                     static_cast<offset_t>(elements),
                     0 /* stream */);

thrust::device_vector<char> temp(temp_size);
auto *temp_storage = thrust::raw_pointer_cast(temp.data());

Finally, we can execute the timed region of the benchmark, which contains the second call to a CUB algorithm and performs the actual work we want to benchmark:

  state.exec(nvbench::exec_tag::no_batch, [&](nvbench::launch &launch) {
    dispatch_t::Dispatch(temp_storage,
                         temp_size,
                         d_in,
                         d_out,
                         static_cast<offset_t>(elements),
                         launch.get_stream());
  });
}

This concludes defining the benchmark function. Now we need to tell NVBench about it:

NVBENCH_BENCH_TYPES(algname, NVBENCH_TYPE_AXES(all_types, offset_types))
  .set_name("base")
  .set_type_axes_names({"T{ct}", "OffsetT{ct}"})
  .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4));

NVBENCH_BENCH_TYPES registers the benchmark as one with multiple compile-time workloads, which are defined by the Cartesian product of the type lists in NVBENCH_TYPE_AXES. set_name(...) sets the name of the benchmark. Only alphabetical characters, numbers and underscores are allowed in the benchmark name.

Furthermore, compile-time axes should be suffixed with {ct}. The runtime axes might be optionally annotated as {io} which stands for importance-ordered. This will tell the tuning infrastructure that the later values on the axis are more important. If the axis is not annotated, each value will be treated as equally important.

When you define a type axis annotated with {ct}, you should consider optimizing the build time. Many variants are going to be build, but the search is considering one compile-time use case at a time. This means that if you have many types to tune for, you’ll end up having many template specializations that you don’t need. To avoid this, for each compile time axis, the tuning framework will predefine a TUNE_AxisName macro with the type that’s currently being tuned. For instance, if you have the type axes T{ct} and OffsetT (as shown above), you can use the following pattern to narrow down the types you compile for:

#ifdef TUNE_T
using all_types = nvbench::type_list<TUNE_T>;
#else
using all_types = nvbench::type_list<char, short, int, long, ...>;
#endif

#ifdef TUNE_OffsetT
using offset_types = nvbench::type_list<TUNE_OffsetT>;
#else
using offset_types = nvbench::type_list<int32_t, int64_t>;
#endif

This logic is already implemented if you use any of the following predefined type lists:

Predefined type lists
Axis name	C++ identifier	Included types
`T{ct}`	`integral_types`	`int8_t, int16_t, int32_t, int64_t`
`T{ct}`	`fundamental_types`	`integral_types` and `int128_t, float, double`
`T{ct}`	`all_types`	`fundamental_types` and `complex`
`OffsetT{ct}`	`offset_types`	`int32_t, int64_t`

But you are free to define your own axis names and use the logic above for them (see the sort pairs example).

A single benchmark file can define multiple benchmarks (multiple benchmark functions registered with NVBENCH_BENCH_TYPES). All benchmarks in a single file must share the same compile-time axes. The tuning infrastructure will run all benchmarks in a single file together for the same compile-time workload and compute a common score across all benchmarks and runtime workloads. This is useful to tune an algorithm for multiple runtime use cases at once, that we don’t intend to provide separate tuning policies for.

Search Process

To get started with tuning, you need to configure CMake. You can use the following command:

$ mkdir build
$ cd build
$ cmake .. --preset=cub-tune -DCMAKE_CUDA_ARCHITECTURES=90 # TODO: Set your GPU architecture

You can then run the tuning search for a specific algorithm and compile-time workload:

$ ../benchmarks/scripts/search.py -R '.*merge_sort.*pairs' -a 'KeyT{ct}=I128' -a 'Elements{io}[pow2]=28'
cub.bench.merge_sort.pairs.trp_0.ld_1.ipt_13.tpb_6 0.6805093269929858
cub.bench.merge_sort.pairs.trp_0.ld_1.ipt_11.tpb_10 1.0774560502969677
...

This will tune merge sort for key-value pairs, for the key type int128_t on 2^28 elements. The -R and -a options are optional. If not specified, all benchmarks are going to be tuned. The -R option can select multiple benchmarks using a regular expression. For the axis option -a, you can also specify a range of values like -a 'KeyT{ct}=[I32,I64]'. Any axis values not supported by a selected benchmark will be ignored. The first variant cub.bench.merge_sort.pairs.trp_0.ld_1.ipt_13.tpb_6 has a score <1 and is thus generally slower than baseline, whereas the second variant cub.bench.merge_sort.pairs.trp_0.ld_1.ipt_11.tpb_10 has a score of >1 and is thus an improvement over the baseline.

Notice there is currently a limitation in search.py which will only execute runs for the first axis value for each axis (independently of whether the axis is specified on the command line or not). Please see this issue for more information.

The tuning framework will handle building the benchmarks (base and variants) by itself. It will keep track of the build time for base and variants. Sometimes, a tuning variant may lead the compiler to hang or take exceptionally long to compile. To keep the tuning process going, if the build time of a variant exceeds a threshold, the build is cancelled. The same applies to benchmarks running for too long.

To get quick feedback on what benchmarks are selected and how big the search space is, you can add the -l option:

$ ../benchmarks/scripts/search.py -R '.*merge_sort.*pairs' -a 'KeyT{ct}=I128' -a 'Elements{io}[pow2]=28' -l
ctk:  12.6.85
cccl:  v2.7.0
### Benchmarks
  * `cub.bench.merge_sort.pairs`: 540 variants:
    * `trp`: (0, 2, 1)
    * `ld`: (0, 3, 1)
    * `ipt`: (7, 25, 1)
    * `tpb`: (6, 11, 1)

It will list all selected benchmarks as well as the total number of variants (the magnitude of the search space) as a result of the Cartesian product of all its tuning parameter spaces.

Analyzing the results

The result of the search is stored in the build/cccl_meta_bench.db file. To analyze the result you can use the analyze.py script. The --coverage flag will show the amount of variants that were covered per compile-time workload:

$ ../benchmarks/scripts/analyze.py --coverage
  cub.bench.radix_sort.keys[T{ct}=I8, OffsetT{ct}=I32] coverage: 167 / 522 (31.9923%)
  cub.bench.radix_sort.keys[T{ct}=I8, OffsetT{ct}=I64] coverage: 152 / 522 (29.1188%)

The --top N flag will list the best N variants for each compile-time workload:

$ ../benchmarks/scripts/analyze.py --top=5
  cub.bench.radix_sort.keys[T{ct}=I8, OffsetT{ct}=I32]:
            variant     score      mins     means      maxs
  97  ipt_19.tpb_512  1.141015  1.039052  1.243448  1.679558
  84  ipt_18.tpb_512  1.136463  1.030434  1.245825  1.668038
  68  ipt_17.tpb_512  1.132696  1.020470  1.250665  1.688889
  41  ipt_15.tpb_576  1.124077  1.011560  1.245011  1.722379
  52  ipt_16.tpb_512  1.121044  0.995238  1.252378  1.717514
  cub.bench.radix_sort.keys[T{ct}=I8, OffsetT{ct}=I64]:
            variant     score      mins     means      maxs
  71  ipt_19.tpb_512  1.250941  1.155738  1.321665  1.647868
  86  ipt_20.tpb_512  1.250840  1.128940  1.308591  1.612382
  55  ipt_17.tpb_512  1.244399  1.152033  1.327424  1.692091
  98  ipt_21.tpb_448  1.231045  1.152798  1.298332  1.621110
  85  ipt_20.tpb_480  1.229382  1.135447  1.294937  1.631225

The name of the variant contains the short parameter names and values used for the variant. For each variant, a score is reported. The base has a score of 1.0, so each score higher than 1.0 is an improvement over the base. However, because a single variant contains multiple runtime workloads, also the minimum, mean, maximum score is reported. If all those three values are larger than 1.0, the variant is strictly better than the base. If only the mean or max are larger than 1.0, the variant may perform better in most runtime workloads, but regress in others. This information can be used to change the existing tuning policies in CUB.

$ ../benchmarks/scripts/analyze.py --variant='ipt_(18|19).tpb_512'

The last command plots distribution of the elapsed times for the specified variants.