Execution Environments#

Most CUB device-wide algorithms accept an optional execution environment as their last argument. The environment is a single object that bundles everything CUB needs to know about how to run the algorithm.

Part of what it carries is familiar: the CUDA stream, which CUB algorithms have always accepted, now simply travels inside the environment. The rest are controls that had no place in the classic API and are enabled by environments: determinism requirements, custom tuning policy selectors, and memory resources. All properties are optional and freely composable with each other.

This page explains what the CUB environment APIs are, and how to use them.

Why environments?#

The classic two-phase CUB API requires three steps: query temporary-storage size, allocate, then execute. That is fine for situations where precise control over the temporary storage allocation is required, or when the temporary storage size needs to be queried independently of allocating it. For example, when the temporary storage is combined with other storage, e.g. for an output, into a single allocation. But in many cases this is not needed and the two-step API is just repetitive boilerplate.

The environment-based single-phase API collapses all of that into one call. The algorithm queries the environment for the properties it needs — for example, the stream to run on, or a memory resource to allocate temporary storage from — and then executes. Properties an algorithm does not use are simply ignored, which is what makes one environment safe to pass to many different algorithms:

cub::DeviceReduce::Sum(d_input, d_output, num_items, env);

Note

The environment argument is entirely optional: since it is defaulted, an algorithm can be invoked with no temporary-storage arguments and no environment at all, and every property falls back to its default (see Default behavior):

cub::DeviceReduce::Sum(d_input, d_output, num_items);

Building an environment #

Some types are already valid environments on their own and can be passed directly to an algorithm. The most common case is a stream: cuda::stream_ref (or a raw cudaStream_t) passed as the last argument is treated as an environment containing just that stream.

auto op     = cuda::std::plus{};
auto input  = thrust::device_vector<float>{0.0f, 1.0f, 2.0f, 3.0f};
auto output = thrust::device_vector<float>(1);
auto init   = 0.0f;

cuda::stream stream{cuda::devices[0]};
cuda::stream_ref stream_ref{stream};

auto error = cub::DeviceReduce::Reduce(input.begin(), output.begin(), input.size(), op, init, stream_ref);
if (error != cudaSuccess)
{
  std::cerr << "cub::DeviceReduce::Reduce failed with status: " << error << '\n';
}

thrust::device_vector<float> expected{6.0f};

When you need more than one property, combine them with cuda::std::execution::env. Properties can be listed in any order. The example below composes a stream, a memory pool for temporary storage, and a determinism requirement (the required declarations are provided by <cuda/execution>, <cuda/stream>, <cuda/memory_pool>, and <cuda/devices>):

// Setup device, stream, memory resource, determinism
auto device          = cuda::devices[device_ordinal];
auto stream          = cuda::stream{device};
auto memory_resource = cuda::device_default_memory_pool(device);
auto determinism     = cuda::execution::require(cuda::execution::determinism::run_to_run);

// Create environment
auto env = cuda::std::execution::env{cuda::stream_ref{stream}, memory_resource, determinism};

// Run
CubDebugExit(cub::DeviceReduce::Sum(d_in.data(), d_out.data(), num_items, env));

The same env object can be passed to multiple algorithm calls without rebuilding it each time, see Reusing an environment across multiple calls.

How to use environments #

The controls below are what environments enable beyond the classic API. A few algorithms additionally accept algorithm-specific controls — e.g. tie-breaking and output-ordering for cub::DeviceTopK. The set of supported properties keeps growing.

Determinism Requirements #

Use cuda::execution::require to request a guarantee, like a reproducible/deterministic execution:

auto op     = cuda::std::plus{};
auto input  = thrust::device_vector<float>{0.0f, 1.0f, 2.0f, 3.0f};
auto output = thrust::device_vector<float>(1);
auto init   = 0.0f;

auto env = cuda::execution::require(cuda::execution::determinism::run_to_run);

auto error = cub::DeviceReduce::Reduce(input.begin(), output.begin(), input.size(), op, init, env);
if (error != cudaSuccess)
{
  std::cerr << "cub::DeviceReduce::Reduce failed with status: " << error << '\n';
}

thrust::device_vector<float> expected{6.0f};

Multiple determinism levels are available. The meaning of each is described in the CCCL determinism overview. Which algorithms support which levels, and each algorithm’s default, are listed in the CUB determinism support matrix. Requesting a level an algorithm does not support is rejected at compile time.

Custom Tuning Policy Selectors #

Pass a custom policy selector through the environment to override CUB’s built-in tuning. A policy selector answers the question “which tuning parameters should this algorithm use on this GPU?”: it is a function object that CUB calls with the GPU’s compute capability, and that returns the tuning parameters — threads per block, items per thread, and so on — packaged as the algorithm’s policy (here cub::MergePolicy):

struct MergePolicySelector
{
  __host__ __device__ constexpr auto operator()(cuda::compute_capability cc) const -> cub::MergePolicy
  {
    return {.threads_per_block        = 512,
            .items_per_thread         = cc > cuda::compute_capability{9, 0} ? 15 : 11,
            .load_modifier            = cub::LOAD_DEFAULT,
            .store_algorithm          = cub::BLOCK_STORE_WARP_TRANSPOSE,
            .use_bulk_copy_for_keys   = true,
            .use_bulk_copy_for_values = false};
  }
};

The selector is wrapped with cuda::execution::tune and passed as (part of) the environment:

auto keys1  = thrust::device_vector<int>{0, 2, 5};
auto keys2  = thrust::device_vector<int>{0, 3, 3, 4};
auto result = thrust::device_vector<int>(7, thrust::no_init);

const auto error = cub::DeviceMerge::MergeKeys(
  keys1.begin(),
  static_cast<int>(keys1.size()),
  keys2.begin(),
  static_cast<int>(keys2.size()),
  result.begin(),
  cuda::std::less{},
  cuda::execution::tune(MergePolicySelector{}));

if (error != cudaSuccess)
{
  std::cerr << "cub::DeviceMerge::MergeKeys failed with status: " << error << '\n';
}

thrust::device_vector<int> expected{0, 0, 2, 3, 3, 4, 5};

Memory Resources #

The memory resource controls where an algorithm’s temporary storage is allocated from. Memory resource types are valid environments on their own, so they can be passed directly or composed with other properties:

auto pool = cuda::device_default_memory_pool(cuda::devices[0]);

// Pass directly...
cub::DeviceReduce::Sum(d_input, d_output, num_items, pool);

// ...or compose with other properties
auto env = cuda::std::execution::env{stream, pool};
cub::DeviceReduce::Sum(d_input, d_output, num_items, env);

Temporary storage is allocated from the memory resource on the algorithm’s stream before execution and released on the same stream afterwards. When no memory resource is present in the environment, CUB falls back to a stream-ordered cudaMallocAsync/cudaFree allocator (see Default behavior).

Reusing an environment across multiple calls #

An env object can be built once and passed to as many algorithm calls as you like:

auto stream = cuda::stream{cuda::devices[0]};
auto pool   = cuda::device_default_memory_pool(cuda::devices[0]);
auto env    = cuda::std::execution::env{cuda::stream_ref{stream}, pool};

cub::DeviceScan::ExclusiveSum(d_a, d_out_a, n, env);
cub::DeviceReduce::Sum(d_b, d_out_b, n, env);
cub::DeviceSelect::If(d_c, d_out_c, d_num_selected, n, my_predicate, env);

stream.sync();

All three calls share the same stream and memory pool. Temporary storage is allocated and released from the pool independently for each call.

Note

When a tuning policy selector is embedded in the environment, it applies to every algorithm that can be tuned by this selector. For example, a policy selector returning a cub::ReducePolicy will be used by all calls to cub::DeviceReduce::* that use the same environment. If two algorithms using the same environment need different tunings, build separate environments.

Default behavior #

The environment argument is optional on every algorithm that supports it. CUB applies the following defaults when a property is absent from the environment (or when no environment is passed at all):

Property	Default when absent
Stream	The default CUDA stream (`cudaStream_t{}`).
Memory resource	A stream-ordered allocator based on `cudaMallocAsync`. Temporary storage is allocated on the stream the algorithm runs on.
Determinism	Algorithm-specific. See Determinism.
Policy selector	CUB’s built-in selector, returning architecture-tuned defaults for the current device.

Missing properties never cause an error. The algorithm simply falls back to its default for that property.