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 Tuning Infrastructure, a set of tools facilitating the process of selecting optimal tuning parameters for a given device and data type.

Definitions

Terms might be ambiguous in a generic context. Below, we omit the word “tuning” but assume it in all definitions. Algorithms are tuned for different workloads. For instance, radix sort can be tuned for different key types, different number of keys, and different distribution of keys. We separate tuning parameters 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.

  • 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.

  • 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 from corresponding search space.

  • Base - a variant that CUB uses by default.

  • Score - a single number representing the performance for a given compile-time workload and 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.

Contributing Benchmarks

There are a few constraints on benchmarks. First of all, all benchmarks in a single file should share type axes. Only alphabetical characters, numbers and underscore are allowed in the benchmark name. The name of the file represents the name of the algorithm. For instance, the benchmarks/bench/radix_sort/keys.cu file name is going to be transformed into cub.bench.radix_sort.keys that’s further used in the infrastructure.

You start writing a benchmark by including nvbench_helper.cuh file. 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 comments. The format consists of %RANGE% keyword, a parameter name, and a range. The range is represented by three numbers: start, end, and step. For instance, the following code defines a search space for 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 docs.

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

Now we have to specialize the dispatch layer. The tuning infrastructure will use TUNE_BASE macro to distinguish between the base and the variant. When base is used, do not specify the policy, so that the default one is used. If the macro is not defined, specify custom policy using macro names defined at the search space specification step.

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

If possible, do not initialize the input data in the benchmark function. Instead, use the gen function. This function 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);

You can optionally add memory usage to the state:

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

Now we are ready to 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<nvbench::uint8_t> temp(temp_size);
auto *temp_storage = thrust::raw_pointer_cast(temp.data());

Finally, we can run the algorithm:

  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());
  });
}

Having the benchmark function, we can tell nvbench about it. A few things to note here. First of all, compile-time axes should be annotated as {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.

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));

When you define a type axis that’s annotated as {ct}, you might want to 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 specializations that you don’t need. To avoid this, for each compile time axis, you can expect a TUNE_AxisName macro with the type that’s currently being tuned. For instance, if you have a type axes T{ct} and OffsetT (as shown above), you can use the following construct:

#ifdef TUNE_T
using types = nvbench::type_list<TUNE_T>;
#else
using types = all_types;
#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 automatically applied to all_types, offset_types, and fundamental_types lists when you use matching names for the axes. You can define your own axis names and use the logic above for them (see sort pairs example).

Search Process

To get started with tuning / benchmarking, you need to configure CMake. The following options are available:

  • CUB_ENABLE_BENCHMARKS - enable bases (default: OFF).

  • CUB_ENABLE_TUNING - enable variants (default: OFF).

Having configured CMake, you can start the search process. Note that the search has to be started from the build directory.

$ cd build
$ cmake -DThrust_DIR=path-to-thrust/thrust/cmake -DCUB_ENABLE_TUNING=YES -DCUB_ENABLE_BENCHMARKS=YES -DCMAKE_BUILD_TYPE=Release -DCMAKE_CUDA_ARCHITECTURES="90" ..
$ ../benchmarks/scripts/search.py -a "T{ct}=[I8,I16]" -R ".*algname.*"

Both -a and -R options are optional. The first one is used to specify types to tune for. The second one is used to specify benchmarks to be tuned. If not specified, all benchmarks are going to be tuned.

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:

$ ../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%)

$ ../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

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

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