[INFRA] Import NVIDIA/CCCL upstream as optimization reference library

CCCL (CUDA C++ Core Libraries) provides:
- CUB: device/block/warp-level GPU primitives (reduce, scan, sort, topk)
- Thrust: high-level parallel algorithms (transform_reduce, sort, scan)
- libcudacxx: CUDA C++ standard library (atomics, barriers, memory)
- cudax: experimental features (memory resources, allocators)
- Tuning policies: per-SM hardware-specific algorithm parameters

Competition optimization vectors mapped to CCCL:
- Output TPS (83% weight): warp_reduce, block_reduce, device_topk
- Input TPS (14% weight): device_scan, block_load, prefetch
- Cache TPS (3% weight): prefix caching strategy patterns
- Memory (0.9 util): pooled/cached/buddy allocators

Source: https://github.com/NVIDIA/cccl (shallow clone, HEAD only)
License: Apache-2.0

cccl_upstream/cub/examples/CMakeLists.txt

@@ -0,0 +1,41 @@
+# Some of the examples include the `cub/test/test_util.h` header, which
+# depends on c2h:
+cccl_get_c2h()
+
+if (NOT CUB_ENABLE_LAUNCH_NO_LAUNCHER)
+  # Examples do not define lid variants; treat them as core-only artifacts.
+  return()
+endif()
+
+## cub_add_example
+#
+# Add an example executable and register it with ctest.
+#
+# target_name_var: Variable name to overwrite with the name of the example
+#   target. Useful for modifying the example/target after creation.
+# example_name: The name of the example minus "cub.example." For
+#   instance, examples/vector.cu will be "vector", and examples/cuda/copy.cu
+#   would be "cuda.copy".
+# example_src: The source file that implements the example.
+#
+function(cub_add_example target_name_var example_name example_src)
+  # The actual name of the test's target:
+  set(example_target cub.example.${example_name})
+  set(${target_name_var} ${example_target} PARENT_SCOPE)
+
+  cccl_add_executable(${example_target} SOURCES "${example_src}" ADD_CTEST)
+  cub_configure_cuda_target(${example_target} RDC ${CUB_FORCE_RDC})
+  target_link_libraries(
+    ${example_target}
+    PRIVATE #
+      cub.compiler_interface
+      cccl.c2h
+  )
+  target_include_directories(
+    ${example_target}
+    PRIVATE "${CUB_SOURCE_DIR}/examples"
+  )
+endfunction()
+
+add_subdirectory(block)
+add_subdirectory(device)

cccl_upstream/cub/examples/block/.gitignore

@@ -0,0 +1,7 @@
+/bin
+/Debug
+/Release
+/cuda55.sdf
+/cuda55.suo
+/cuda60.sdf
+/cuda60.suo

cccl_upstream/cub/examples/block/CMakeLists.txt

@@ -0,0 +1,18 @@
+file(
+  GLOB_RECURSE example_srcs
+  RELATIVE "${CMAKE_CURRENT_LIST_DIR}"
+  CONFIGURE_DEPENDS
+  example_*.cu
+)
+
+foreach (example_src IN LISTS example_srcs)
+  get_filename_component(example_name "${example_src}" NAME_WE)
+  string(
+    REGEX REPLACE
+    "^example_block_"
+    "block."
+    example_name
+    "${example_name}"
+  )
+  cub_add_example(target_name ${example_name} "${example_src}")
+endforeach()

cccl_upstream/cub/examples/block/example_block_radix_sort.cu

@@ -0,0 +1,315 @@
+// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
+// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+/******************************************************************************
+ * Simple demonstration of cub::BlockRadixSort
+ *
+ * To compile using the command line:
+ *   nvcc -arch=sm_XX example_block_radix_sort.cu -I../.. -lcudart -O3
+ *
+ ******************************************************************************/
+
+// Ensure printing of CUDA runtime errors to console (define before including cub.h)
+#define CUB_STDERR
+
+#include <cub/block/block_load.cuh>
+#include <cub/block/block_radix_sort.cuh>
+#include <cub/block/block_store.cuh>
+
+#include <algorithm>
+#include <cstdio>
+#include <iostream>
+
+#include "../../test/test_util.h"
+
+using namespace cub;
+
+//---------------------------------------------------------------------
+// Globals, constants and aliases
+//---------------------------------------------------------------------
+
+/// Verbose output
+bool g_verbose = false;
+
+/// Timing iterations
+int g_timing_iterations = 100;
+
+/// Default grid size
+int g_grid_size = 1;
+
+/// Uniform key samples
+bool g_uniform_keys;
+
+//---------------------------------------------------------------------
+// Kernels
+//---------------------------------------------------------------------
+
+/**
+ * Simple kernel for performing a block-wide sorting over integers
+ */
+template <typename Key,
+          int BLOCK_THREADS,
+          int ITEMS_PER_THREAD>
+__launch_bounds__(BLOCK_THREADS) __global__
+  void BlockSortKernel(Key* d_in, // Tile of input
+                       Key* d_out, // Tile of output
+                       clock_t* d_elapsed) // Elapsed cycle count of block scan
+{
+  static constexpr int TILE_SIZE = BLOCK_THREADS * ITEMS_PER_THREAD;
+
+  // Specialize BlockLoad type for our thread block (uses warp-striped loads for coalescing, then transposes in shared
+  // memory to a blocked arrangement)
+  using BlockLoadT = BlockLoad<Key, BLOCK_THREADS, ITEMS_PER_THREAD, BLOCK_LOAD_WARP_TRANSPOSE>;
+
+  // Specialize BlockRadixSort type for our thread block
+  using BlockRadixSortT = BlockRadixSort<Key, BLOCK_THREADS, ITEMS_PER_THREAD>;
+
+  // Shared memory
+  __shared__ union TempStorage
+  {
+    typename BlockLoadT::TempStorage load;
+    typename BlockRadixSortT::TempStorage sort;
+  } temp_storage;
+
+  // Per-thread tile items
+  Key items[ITEMS_PER_THREAD];
+
+  // Our current block's offset
+  int block_offset = blockIdx.x * TILE_SIZE;
+
+  // Load items into a blocked arrangement
+  BlockLoadT(temp_storage.load).Load(d_in + block_offset, items);
+
+  // Barrier for smem reuse
+  __syncthreads();
+
+  // Start cycle timer
+  clock_t start = clock();
+
+  // Sort keys
+  BlockRadixSortT(temp_storage.sort).SortBlockedToStriped(items);
+
+  // Stop cycle timer
+  clock_t stop = clock();
+
+  // Store output in striped fashion
+  StoreDirectStriped<BLOCK_THREADS>(threadIdx.x, d_out + block_offset, items);
+
+  // Store elapsed clocks
+  if (threadIdx.x == 0)
+  {
+    d_elapsed[blockIdx.x] = (start > stop) ? start - stop : stop - start;
+  }
+}
+
+//---------------------------------------------------------------------
+// Host utilities
+//---------------------------------------------------------------------
+
+/**
+ * Initialize sorting problem (and solution).
+ */
+template <typename Key>
+void Initialize(Key* h_in, Key* h_reference, int num_items, int tile_size)
+{
+  for (int i = 0; i < num_items; ++i)
+  {
+    if (g_uniform_keys)
+    {
+      h_in[i] = 0;
+    }
+    else
+    {
+      RandomBits(h_in[i]);
+    }
+    h_reference[i] = h_in[i];
+  }
+
+  // Only sort the first tile
+  std::sort(h_reference, h_reference + tile_size);
+}
+
+/**
+ * Test BlockScan
+ */
+template <typename Key, int BLOCK_THREADS, int ITEMS_PER_THREAD>
+void Test()
+{
+  constexpr int TILE_SIZE = BLOCK_THREADS * ITEMS_PER_THREAD;
+
+  // Allocate host arrays
+  Key* h_in          = new Key[TILE_SIZE * g_grid_size];
+  Key* h_reference   = new Key[TILE_SIZE * g_grid_size];
+  clock_t* h_elapsed = new clock_t[g_grid_size];
+
+  // Initialize problem and reference output on host
+  Initialize(h_in, h_reference, TILE_SIZE * g_grid_size, TILE_SIZE);
+
+  // Initialize device arrays
+  Key* d_in          = nullptr;
+  Key* d_out         = nullptr;
+  clock_t* d_elapsed = nullptr;
+  CubDebugExit(cudaMalloc((void**) &d_in, sizeof(Key) * TILE_SIZE * g_grid_size));
+  CubDebugExit(cudaMalloc((void**) &d_out, sizeof(Key) * TILE_SIZE * g_grid_size));
+  CubDebugExit(cudaMalloc((void**) &d_elapsed, sizeof(clock_t) * g_grid_size));
+
+  // Display input problem data
+  if (g_verbose)
+  {
+    printf("Input data: ");
+    for (int i = 0; i < TILE_SIZE; i++)
+    {
+      std::cout << h_in[i] << ", "; // NOLINT(bugprone-unintended-char-ostream-output)
+    }
+    printf("\n\n");
+  }
+
+  // Kernel props
+  int max_sm_occupancy;
+  CubDebugExit(MaxSmOccupancy(max_sm_occupancy, BlockSortKernel<Key, BLOCK_THREADS, ITEMS_PER_THREAD>, BLOCK_THREADS));
+
+  // Copy problem to device
+  CubDebugExit(cudaMemcpy(d_in, h_in, sizeof(Key) * TILE_SIZE * g_grid_size, cudaMemcpyHostToDevice));
+
+  printf(
+    "BlockRadixSort %d items (%d timing iterations, %d blocks, %d threads, %d items per thread, %d SM occupancy):\n",
+    TILE_SIZE * g_grid_size,
+    g_timing_iterations,
+    g_grid_size,
+    BLOCK_THREADS,
+    ITEMS_PER_THREAD,
+    max_sm_occupancy);
+  fflush(stdout);
+
+  // Run kernel once to prime caches and check result
+  BlockSortKernel<Key, BLOCK_THREADS, ITEMS_PER_THREAD><<<g_grid_size, BLOCK_THREADS>>>(d_in, d_out, d_elapsed);
+
+  // Check for kernel errors and STDIO from the kernel, if any
+  CubDebugExit(cudaPeekAtLastError());
+  CubDebugExit(cudaDeviceSynchronize());
+
+  // Check results
+  printf("\tOutput items: ");
+  int compare = CompareDeviceResults(h_reference, d_out, TILE_SIZE, g_verbose, g_verbose);
+  printf("%s\n", compare ? "FAIL" : "PASS");
+  AssertEquals(0, compare);
+  fflush(stdout);
+
+  // Run this several times and average the performance results
+  GpuTimer timer;
+  float elapsed_millis              = 0.0;
+  unsigned long long elapsed_clocks = 0;
+
+  for (int i = 0; i < g_timing_iterations; ++i)
+  {
+    timer.Start();
+
+    // Run kernel
+    BlockSortKernel<Key, BLOCK_THREADS, ITEMS_PER_THREAD><<<g_grid_size, BLOCK_THREADS>>>(d_in, d_out, d_elapsed);
+
+    timer.Stop();
+    elapsed_millis += timer.ElapsedMillis();
+
+    // Copy clocks from device
+    CubDebugExit(cudaMemcpy(h_elapsed, d_elapsed, sizeof(clock_t) * g_grid_size, cudaMemcpyDeviceToHost));
+    for (int j = 0; j < g_grid_size; j++)
+    {
+      elapsed_clocks += h_elapsed[j];
+    }
+  }
+
+  // Check for kernel errors and STDIO from the kernel, if any
+  CubDebugExit(cudaDeviceSynchronize());
+
+  // Display timing results
+  float avg_millis           = elapsed_millis / static_cast<float>(g_timing_iterations);
+  float avg_items_per_sec    = float(TILE_SIZE * g_grid_size) / avg_millis / 1000.0f;
+  double avg_clocks          = double(elapsed_clocks) / g_timing_iterations / g_grid_size;
+  double avg_clocks_per_item = avg_clocks / TILE_SIZE;
+
+  printf("\tAverage BlockRadixSort::SortBlocked clocks: %.3f\n", avg_clocks);
+  printf("\tAverage BlockRadixSort::SortBlocked clocks per item: %.3f\n", avg_clocks_per_item);
+  printf("\tAverage kernel millis: %.4f\n", avg_millis);
+  printf("\tAverage million items / sec: %.4f\n", avg_items_per_sec);
+  fflush(stdout);
+
+  // Cleanup
+  if (h_in)
+  {
+    delete[] h_in;
+  }
+  if (h_reference)
+  {
+    delete[] h_reference;
+  }
+  if (h_elapsed)
+  {
+    delete[] h_elapsed;
+  }
+  if (d_in)
+  {
+    CubDebugExit(cudaFree(d_in));
+  }
+  if (d_out)
+  {
+    CubDebugExit(cudaFree(d_out));
+  }
+  if (d_elapsed)
+  {
+    CubDebugExit(cudaFree(d_elapsed));
+  }
+}
+
+/**
+ * Main
+ */
+int main(int argc, char** argv)
+{
+  // Initialize command line
+  CommandLineArgs args(argc, argv);
+  g_verbose      = args.CheckCmdLineFlag("v");
+  g_uniform_keys = args.CheckCmdLineFlag("uniform");
+  args.GetCmdLineArgument("i", g_timing_iterations);
+  args.GetCmdLineArgument("grid-size", g_grid_size);
+
+  // Print usage
+  if (args.CheckCmdLineFlag("help"))
+  {
+    printf("%s "
+           "[--device=<device-id>] "
+           "[--i=<timing iterations (default:%d)>]"
+           "[--grid-size=<grid size (default:%d)>]"
+           "[--v] "
+           "\n",
+           argv[0],
+           g_timing_iterations,
+           g_grid_size);
+    exit(0);
+  }
+
+  // Initialize device
+  CubDebugExit(args.DeviceInit());
+  fflush(stdout);
+
+  // Run tests
+  printf("\nuint32:\n");
+  fflush(stdout);
+  Test<unsigned int, 128, 13>();
+  printf("\n");
+  fflush(stdout);
+
+  printf("\nfp32:\n");
+  fflush(stdout);
+  Test<float, 128, 13>();
+  printf("\n");
+  fflush(stdout);
+
+  printf("\nuint8:\n");
+  fflush(stdout);
+  Test<unsigned char, 128, 13>();
+  printf("\n");
+  fflush(stdout);
+
+  return 0;
+}

cccl_upstream/cub/examples/block/example_block_reduce.cu

@@ -0,0 +1,273 @@
+// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
+// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+/******************************************************************************
+ * Simple demonstration of cub::BlockReduce
+ *
+ * To compile using the command line:
+ *   nvcc -arch=sm_XX example_block_reduce.cu -I../.. -lcudart -O3
+ *
+ ******************************************************************************/
+
+// Ensure printing of CUDA runtime errors to console (define before including cub.h)
+#define CUB_STDERR
+
+#include <cub/block/block_load.cuh>
+#include <cub/block/block_reduce.cuh>
+#include <cub/block/block_store.cuh>
+
+#include <cstdio>
+#include <iostream>
+
+#include "../../test/test_util.h"
+
+using namespace cub;
+
+//---------------------------------------------------------------------
+// Globals, constants and aliases
+//---------------------------------------------------------------------
+
+/// Verbose output
+bool g_verbose = false;
+
+/// Timing iterations
+int g_timing_iterations = 100;
+
+/// Default grid size
+int g_grid_size = 1;
+
+//---------------------------------------------------------------------
+// Kernels
+//---------------------------------------------------------------------
+
+/**
+ * Simple kernel for performing a block-wide reduction.
+ */
+template <int BLOCK_THREADS,
+          int ITEMS_PER_THREAD,
+          BlockReduceAlgorithm ALGORITHM>
+__global__ void BlockReduceKernel(int* d_in, // Tile of input
+                                  int* d_out, // Tile aggregate
+                                  clock_t* d_elapsed) // Elapsed cycle count of block reduction
+{
+  // Specialize BlockReduce type for our thread block
+  using BlockReduceT = BlockReduce<int, BLOCK_THREADS, ALGORITHM>;
+
+  // Shared memory
+  __shared__ typename BlockReduceT::TempStorage temp_storage;
+
+  // Per-thread tile data
+  int data[ITEMS_PER_THREAD];
+  LoadDirectStriped<BLOCK_THREADS>(threadIdx.x, d_in, data);
+
+  // Start cycle timer
+  clock_t start = clock();
+
+  // Compute sum
+  int aggregate = BlockReduceT(temp_storage).Sum(data);
+
+  // Stop cycle timer
+  clock_t stop = clock();
+
+  // Store aggregate and elapsed clocks
+  if (threadIdx.x == 0)
+  {
+    *d_elapsed = (start > stop) ? start - stop : stop - start;
+    *d_out     = aggregate;
+  }
+}
+
+//---------------------------------------------------------------------
+// Host utilities
+//---------------------------------------------------------------------
+
+/**
+ * Initialize reduction problem (and solution).
+ * Returns the aggregate
+ */
+int Initialize(int* h_in, int num_items)
+{
+  int inclusive = 0;
+
+  for (int i = 0; i < num_items; ++i)
+  {
+    h_in[i] = i % 17;
+    inclusive += h_in[i];
+  }
+
+  return inclusive;
+}
+
+/**
+ * Test thread block reduction
+ */
+template <int BLOCK_THREADS, int ITEMS_PER_THREAD, BlockReduceAlgorithm ALGORITHM>
+void Test()
+{
+  constexpr int TILE_SIZE = BLOCK_THREADS * ITEMS_PER_THREAD;
+
+  // Allocate host arrays
+  int* h_in  = new int[TILE_SIZE];
+  int* h_gpu = new int[TILE_SIZE + 1];
+
+  // Initialize problem and reference output on host
+  int h_aggregate = Initialize(h_in, TILE_SIZE);
+
+  // Initialize device arrays
+  int* d_in          = nullptr;
+  int* d_out         = nullptr;
+  clock_t* d_elapsed = nullptr;
+  cudaMalloc((void**) &d_in, sizeof(int) * TILE_SIZE);
+  cudaMalloc((void**) &d_out, sizeof(int) * 1);
+  cudaMalloc((void**) &d_elapsed, sizeof(clock_t));
+
+  // Display input problem data
+  if (g_verbose)
+  {
+    printf("Input data: ");
+    for (int i = 0; i < TILE_SIZE; i++)
+    {
+      printf("%d, ", h_in[i]);
+    }
+    printf("\n\n");
+  }
+
+  // Kernel props
+  int max_sm_occupancy;
+  CubDebugExit(
+    MaxSmOccupancy(max_sm_occupancy, BlockReduceKernel<BLOCK_THREADS, ITEMS_PER_THREAD, ALGORITHM>, BLOCK_THREADS));
+
+  // Copy problem to device
+  cudaMemcpy(d_in, h_in, sizeof(int) * TILE_SIZE, cudaMemcpyHostToDevice);
+
+  printf("BlockReduce algorithm %s on %d items (%d timing iterations, %d blocks, %d threads, %d items per thread, %d "
+         "SM occupancy):\n",
+         (ALGORITHM == BLOCK_REDUCE_RAKING) ? "BLOCK_REDUCE_RAKING" : "BLOCK_REDUCE_WARP_REDUCTIONS",
+         TILE_SIZE,
+         g_timing_iterations,
+         g_grid_size,
+         BLOCK_THREADS,
+         ITEMS_PER_THREAD,
+         max_sm_occupancy);
+
+  // Run kernel
+  BlockReduceKernel<BLOCK_THREADS, ITEMS_PER_THREAD, ALGORITHM><<<g_grid_size, BLOCK_THREADS>>>(d_in, d_out, d_elapsed);
+
+  // Check total aggregate
+  printf("\tAggregate: ");
+  int compare = CompareDeviceResults(&h_aggregate, d_out, 1, g_verbose, g_verbose);
+  printf("%s\n", compare ? "FAIL" : "PASS");
+  AssertEquals(0, compare);
+
+  // Run this several times and average the performance results
+  GpuTimer timer;
+  float elapsed_millis   = 0.0;
+  clock_t elapsed_clocks = 0;
+
+  for (int i = 0; i < g_timing_iterations; ++i)
+  {
+    // Copy problem to device
+    cudaMemcpy(d_in, h_in, sizeof(int) * TILE_SIZE, cudaMemcpyHostToDevice);
+
+    timer.Start();
+
+    // Run kernel
+    BlockReduceKernel<BLOCK_THREADS, ITEMS_PER_THREAD, ALGORITHM>
+      <<<g_grid_size, BLOCK_THREADS>>>(d_in, d_out, d_elapsed);
+
+    timer.Stop();
+    elapsed_millis += timer.ElapsedMillis();
+
+    // Copy clocks from device
+    clock_t clocks;
+    CubDebugExit(cudaMemcpy(&clocks, d_elapsed, sizeof(clock_t), cudaMemcpyDeviceToHost));
+    elapsed_clocks += clocks;
+  }
+
+  // Check for kernel errors and STDIO from the kernel, if any
+  CubDebugExit(cudaPeekAtLastError());
+  CubDebugExit(cudaDeviceSynchronize());
+
+  // Display timing results
+  float avg_millis          = elapsed_millis / static_cast<float>(g_timing_iterations);
+  float avg_items_per_sec   = float(TILE_SIZE * g_grid_size) / avg_millis / 1000.0f;
+  float avg_clocks          = float(elapsed_clocks) / static_cast<float>(g_timing_iterations);
+  float avg_clocks_per_item = avg_clocks / TILE_SIZE;
+
+  printf("\tAverage BlockReduce::Sum clocks: %.3f\n", avg_clocks);
+  printf("\tAverage BlockReduce::Sum clocks per item: %.3f\n", avg_clocks_per_item);
+  printf("\tAverage kernel millis: %.4f\n", avg_millis);
+  printf("\tAverage million items / sec: %.4f\n", avg_items_per_sec);
+
+  // Cleanup
+  if (h_in)
+  {
+    delete[] h_in;
+  }
+  if (h_gpu)
+  {
+    delete[] h_gpu;
+  }
+  if (d_in)
+  {
+    cudaFree(d_in);
+  }
+  if (d_out)
+  {
+    cudaFree(d_out);
+  }
+  if (d_elapsed)
+  {
+    cudaFree(d_elapsed);
+  }
+}
+
+/**
+ * Main
+ */
+int main(int argc, char** argv)
+{
+  // Initialize command line
+  CommandLineArgs args(argc, argv);
+  g_verbose = args.CheckCmdLineFlag("v");
+  args.GetCmdLineArgument("i", g_timing_iterations);
+  args.GetCmdLineArgument("grid-size", g_grid_size);
+
+  // Print usage
+  if (args.CheckCmdLineFlag("help"))
+  {
+    printf("%s "
+           "[--device=<device-id>] "
+           "[--i=<timing iterations>] "
+           "[--grid-size=<grid size>] "
+           "[--v] "
+           "\n",
+           argv[0]);
+    exit(0);
+  }
+
+  // Initialize device
+  CubDebugExit(args.DeviceInit());
+
+  // Run tests
+  Test<1024, 1, BLOCK_REDUCE_RAKING>();
+  Test<512, 2, BLOCK_REDUCE_RAKING>();
+  Test<256, 4, BLOCK_REDUCE_RAKING>();
+  Test<128, 8, BLOCK_REDUCE_RAKING>();
+  Test<64, 16, BLOCK_REDUCE_RAKING>();
+  Test<32, 32, BLOCK_REDUCE_RAKING>();
+  Test<16, 64, BLOCK_REDUCE_RAKING>();
+
+  printf("-------------\n");
+
+  Test<1024, 1, BLOCK_REDUCE_WARP_REDUCTIONS>();
+  Test<512, 2, BLOCK_REDUCE_WARP_REDUCTIONS>();
+  Test<256, 4, BLOCK_REDUCE_WARP_REDUCTIONS>();
+  Test<128, 8, BLOCK_REDUCE_WARP_REDUCTIONS>();
+  Test<64, 16, BLOCK_REDUCE_WARP_REDUCTIONS>();
+  Test<32, 32, BLOCK_REDUCE_WARP_REDUCTIONS>();
+  Test<16, 64, BLOCK_REDUCE_WARP_REDUCTIONS>();
+
+  return 0;
+}

cccl_upstream/cub/examples/block/example_block_reduce_dyn_smem.cu

@@ -0,0 +1,215 @@
+// SPDX-FileCopyrightText: Copyright (c) 2011-2021, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+/******************************************************************************
+ * Simple demonstration of cub::BlockReduce with dynamic shared memory
+ *
+ * To compile using the command line:
+ *   nvcc -arch=sm_XX example_block_reduce_dyn_smem.cu -I../.. -lcudart -O3 -std=c++17
+ *
+ ******************************************************************************/
+
+// Ensure printing of CUDA runtime errors to console (define before including cub.h)
+#define CUB_STDERR
+
+#include <cub/block/block_load.cuh>
+#include <cub/block/block_reduce.cuh>
+#include <cub/block/block_store.cuh>
+
+#include <algorithm>
+#include <cstdio>
+#include <iostream>
+
+#include "../../test/test_util.h"
+
+using namespace cub;
+
+//---------------------------------------------------------------------
+// Globals, constants and aliases
+//---------------------------------------------------------------------
+
+/// Verbose output
+bool g_verbose = false;
+
+/// Default grid size
+int g_grid_size = 1;
+
+//---------------------------------------------------------------------
+// Kernels
+//---------------------------------------------------------------------
+
+/**
+ * Simple kernel for performing a block-wide reduction.
+ */
+template <int BLOCK_THREADS>
+__global__ void BlockReduceKernel(int* d_in, // Tile of input
+                                  int* d_out // Tile aggregate
+)
+{
+  // Specialize BlockReduce type for our thread block
+  using BlockReduceT = cub::BlockReduce<int, BLOCK_THREADS>;
+  using TempStorageT = typename BlockReduceT::TempStorage;
+
+  union ShmemLayout
+  {
+    TempStorageT reduce;
+    int aggregate;
+  };
+
+  // shared memory byte-array
+  extern __shared__ __align__(alignof(ShmemLayout)) char smem[];
+
+  // cast to lvalue reference of expected type
+  auto& temp_storage = reinterpret_cast<TempStorageT&>(smem);
+
+  int data = d_in[threadIdx.x];
+
+  // Compute sum
+  int aggregate = BlockReduceT(temp_storage).Sum(data);
+
+  // block-wide sync barrier necessary to re-use shared mem safely
+  __syncthreads();
+  int* smem_integers = reinterpret_cast<int*>(smem);
+  if (threadIdx.x == 0)
+  {
+    smem_integers[0] = aggregate;
+  }
+
+  // sync to make new shared value available to all threads
+  __syncthreads();
+  aggregate = smem_integers[0];
+
+  // all threads write the aggregate to output
+  d_out[threadIdx.x] = aggregate;
+}
+
+//---------------------------------------------------------------------
+// Host utilities
+//---------------------------------------------------------------------
+
+/**
+ * Initialize reduction problem (and solution).
+ * Returns the aggregate
+ */
+int Initialize(int* h_in, int num_items)
+{
+  int inclusive = 0;
+
+  for (int i = 0; i < num_items; ++i)
+  {
+    h_in[i] = i % 17;
+    inclusive += h_in[i];
+  }
+
+  return inclusive;
+}
+
+/**
+ * Test thread block reduction
+ */
+template <int BLOCK_THREADS>
+void Test()
+{
+  // Allocate host arrays
+  int* h_in = new int[BLOCK_THREADS];
+
+  // Initialize problem and reference output on host
+  int h_aggregate = Initialize(h_in, BLOCK_THREADS);
+
+  // Initialize device arrays
+  int* d_in  = nullptr;
+  int* d_out = nullptr;
+  cudaMalloc((void**) &d_in, sizeof(int) * BLOCK_THREADS);
+  cudaMalloc((void**) &d_out, sizeof(int) * BLOCK_THREADS);
+
+  // Display input problem data
+  if (g_verbose)
+  {
+    printf("Input data: ");
+    for (int i = 0; i < BLOCK_THREADS; i++)
+    {
+      printf("%d, ", h_in[i]);
+    }
+    printf("\n\n");
+  }
+
+  // Copy problem to device
+  cudaMemcpy(d_in, h_in, sizeof(int) * BLOCK_THREADS, cudaMemcpyHostToDevice);
+
+  // determine necessary storage size:
+  auto block_reduce_temp_bytes = sizeof(typename cub::BlockReduce<int, BLOCK_THREADS>::TempStorage);
+  // finally, we need to make sure that we can hold at least one integer
+  // needed in the kernel to exchange data after reduction
+  auto smem_size = (std::max) (1 * sizeof(int), block_reduce_temp_bytes);
+
+  // use default stream
+  cudaStream_t stream = nullptr;
+
+  // Run reduction kernel
+  BlockReduceKernel<BLOCK_THREADS><<<g_grid_size, BLOCK_THREADS, smem_size, stream>>>(d_in, d_out);
+
+  // Check total aggregate
+  printf("\tAggregate: ");
+  int compare = 0;
+  for (int i = 0; i < BLOCK_THREADS; i++)
+  {
+    compare = compare || CompareDeviceResults(&h_aggregate, d_out + i, 1, g_verbose, g_verbose);
+  }
+  printf("%s\n", compare ? "FAIL" : "PASS");
+  AssertEquals(0, compare);
+
+  // Check for kernel errors and STDIO from the kernel, if any
+  CubDebugExit(cudaPeekAtLastError());
+  CubDebugExit(cudaDeviceSynchronize());
+
+  // Cleanup
+  if (h_in)
+  {
+    delete[] h_in;
+  }
+  if (d_in)
+  {
+    cudaFree(d_in);
+  }
+  if (d_out)
+  {
+    cudaFree(d_out);
+  }
+}
+
+/**
+ * Main
+ */
+int main(int argc, char** argv)
+{
+  // Initialize command line
+  CommandLineArgs args(argc, argv);
+  g_verbose = args.CheckCmdLineFlag("v");
+  args.GetCmdLineArgument("grid-size", g_grid_size);
+
+  // Print usage
+  if (args.CheckCmdLineFlag("help"))
+  {
+    printf("%s "
+           "[--device=<device-id>] "
+           "[--grid-size=<grid size>] "
+           "[--v] "
+           "\n",
+           argv[0]);
+    exit(0);
+  }
+
+  // Initialize device
+  CubDebugExit(args.DeviceInit());
+
+  // Run tests
+  Test<1024>();
+  Test<512>();
+  Test<256>();
+  Test<128>();
+  Test<64>();
+  Test<32>();
+  Test<16>();
+
+  return 0;
+}

cccl_upstream/cub/examples/device/.gitignore

@@ -0,0 +1,8 @@
+/bin
+/Debug
+/ipch
+/Release
+/cuda55.sdf
+/cuda55.suo
+/cuda60.sdf
+/cuda60.suo

cccl_upstream/cub/examples/device/CMakeLists.txt

@@ -0,0 +1,18 @@
+file(
+  GLOB_RECURSE example_srcs
+  RELATIVE "${CMAKE_CURRENT_LIST_DIR}"
+  CONFIGURE_DEPENDS
+  example_*.cu
+)
+
+foreach (example_src IN LISTS example_srcs)
+  get_filename_component(example_name "${example_src}" NAME_WE)
+  string(
+    REGEX REPLACE
+    "^example_device_"
+    "device."
+    example_name
+    "${example_name}"
+  )
+  cub_add_example(target_name ${example_name} "${example_src}")
+endforeach()

cccl_upstream/cub/examples/device/example_device_decoupled_look_back.cu

@@ -0,0 +1,104 @@
+// SPDX-FileCopyrightText: Copyright (c) 2011-2023, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+#include <cub/device/device_scan.cuh>
+
+#include <thrust/device_vector.h>
+
+#include <iostream>
+
+template <class ScanTileStateT>
+__global__ void init_kernel(ScanTileStateT tile_state, int blocks_in_grid)
+{
+  tile_state.InitializeStatus(blocks_in_grid);
+}
+
+template <class MessageT>
+__global__ void decoupled_look_back_kernel(cub::ScanTileState<MessageT> tile_state)
+{
+  using scan_op_t         = ::cuda::std::plus<>;
+  using scan_tile_state_t = cub::ScanTileState<MessageT>;
+  using tile_prefix_op    = cub::TilePrefixCallbackOp<MessageT, scan_op_t, scan_tile_state_t>;
+  using temp_storage_t    = typename tile_prefix_op::TempStorage;
+
+  // Allocate temp storage in shared memory
+  __shared__ temp_storage_t temp_storage;
+
+  scan_op_t scan_op{};
+  constexpr unsigned int threads_in_warp = 32;
+  const unsigned int tid                 = threadIdx.x;
+
+  // Construct prefix op
+  tile_prefix_op prefix(tile_state, temp_storage, scan_op);
+  const unsigned int tile_idx = prefix.GetTileIdx();
+
+  // Compute block aggregate
+  MessageT block_aggregate = blockIdx.x;
+
+  if (tile_idx == 0)
+  {
+    // There are no blocks to look back to, immediately set the inclusive state
+    if (tid == 0)
+    {
+      tile_state.SetInclusive(tile_idx, block_aggregate);
+      printf("tile %d: inclusive = %d\n", tile_idx, block_aggregate);
+    }
+  }
+  else
+  {
+    // Only the first warp in the block can perform the look back
+    const unsigned int warp_id = tid / threads_in_warp;
+
+    if (warp_id == 0)
+    {
+      // Perform the decoupled look-back.
+      // 1. Publish the block's local aggregate to global memory immediately.
+      //    This allows downstream blocks to include this tile's contribution without waiting for this block to fully
+      //    resolve its global prefix.
+      // 2. Block and traverse predecessor tiles (look-back) to compute the global exclusive prefix for this tile.
+      // 3. Update this tile's global state to 'Prefix' (inclusive sum), creating a checkpoint that stops the look-back
+      // of downstream blocks.
+      // Note, the invocation of the prefix will block until the look-back is complete.
+      MessageT exclusive_prefix = prefix(block_aggregate);
+
+      if (tid == 0)
+      {
+        MessageT inclusive_prefix = scan_op(exclusive_prefix, block_aggregate);
+        printf("tile %d: exclusive = %d inclusive = %d\n", tile_idx, exclusive_prefix, inclusive_prefix);
+      }
+    }
+  }
+}
+
+template <class MessageT>
+void decoupled_look_back_example(int blocks_in_grid)
+{
+  using scan_tile_state_t = cub::ScanTileState<MessageT>;
+
+  // Query temporary storage requirements
+  std::size_t temp_storage_bytes{};
+  scan_tile_state_t::AllocationSize(blocks_in_grid, temp_storage_bytes);
+
+  // Allocate temporary storage
+  thrust::device_vector<std::uint8_t> temp_storage(temp_storage_bytes);
+  std::uint8_t* d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());
+
+  // Initialize temporary storage
+  scan_tile_state_t tile_status;
+  tile_status.Init(blocks_in_grid, d_temp_storage, temp_storage_bytes);
+  constexpr unsigned int threads_in_init_block = 256;
+  const unsigned int blocks_in_init_grid       = ::cuda::ceil_div(blocks_in_grid, threads_in_init_block);
+  init_kernel<<<blocks_in_init_grid, threads_in_init_block>>>(tile_status, blocks_in_grid);
+
+  // Launch decoupled look-back
+  constexpr unsigned int threads_in_block = 256;
+  decoupled_look_back_kernel<<<blocks_in_grid, threads_in_block>>>(tile_status);
+
+  // Wait for kernel to finish
+  cudaDeviceSynchronize();
+}
+
+int main()
+{
+  decoupled_look_back_example<int>(14);
+}

cccl_upstream/cub/examples/device/example_device_partition_flagged.cu

@@ -0,0 +1,222 @@
+// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
+// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+/******************************************************************************
+ * Simple example of DevicePartition::Flagged().
+ *
+ * Partition flagged items from from a sequence of int keys using a
+ * corresponding sequence of unsigned char flags.
+ *
+ * To compile using the command line:
+ *   nvcc -arch=sm_XX example_device_partition_flagged.cu -I../.. -lcudart -O3
+ *
+ ******************************************************************************/
+
+// Ensure printing of CUDA runtime errors to console
+#define CUB_STDERR
+
+#include <cub/device/device_partition.cuh>
+#include <cub/util_allocator.cuh>
+
+#include <cuda/std/limits>
+
+#include <cstdio>
+
+#include "../../test/test_util.h"
+
+using namespace cub;
+
+//---------------------------------------------------------------------
+// Globals, constants and aliases
+//---------------------------------------------------------------------
+
+bool g_verbose = false; // Whether to display input/output to console
+CachingDeviceAllocator g_allocator(true); // Caching allocator for device memory
+
+//---------------------------------------------------------------------
+// Test generation
+//---------------------------------------------------------------------
+
+/**
+ * Initialize problem, setting flags at distances of random length
+ * chosen from [1..max_segment]
+ */
+void Initialize(int* h_in, unsigned char* h_flags, int num_items, int max_segment)
+{
+  int key = 0;
+  int i   = 0;
+  while (i < num_items)
+  {
+    // Select number of repeating occurrences
+    unsigned short bits;
+    RandomBits(bits);
+    const int repeat = cuda::std::max(
+      1,
+      static_cast<int>(static_cast<float>(bits)
+                       * (static_cast<float>(max_segment) / cuda::std::numeric_limits<unsigned short>::max())));
+
+    int j = i;
+    while (j < cuda::std::min(i + repeat, num_items))
+    {
+      h_flags[j] = 0;
+      h_in[j]    = key;
+      j++;
+    }
+
+    h_flags[i] = 1;
+    i          = j;
+    key++;
+  }
+
+  if (g_verbose)
+  {
+    printf("Input:\n");
+    DisplayResults(h_in, num_items);
+    printf("Flags:\n");
+    DisplayResults(h_flags, num_items);
+    printf("\n\n");
+  }
+}
+
+/**
+ * Solve unique problem
+ */
+int Solve(int* h_in, unsigned char* h_flags, int* h_reference, int num_items)
+{
+  int num_selected = 0;
+  for (int i = 0; i < num_items; ++i)
+  {
+    if (h_flags[i])
+    {
+      h_reference[num_selected] = h_in[i];
+      num_selected++;
+    }
+    else
+    {
+      h_reference[num_items - (i - num_selected) - 1] = h_in[i];
+    }
+  }
+
+  return num_selected;
+}
+
+//---------------------------------------------------------------------
+// Main
+//---------------------------------------------------------------------
+
+/**
+ * Main
+ */
+int main(int argc, char** argv)
+{
+  int num_items   = 150;
+  int max_segment = 40; // Maximum segment length
+
+  // Initialize command line
+  CommandLineArgs args(argc, argv);
+  g_verbose = args.CheckCmdLineFlag("v");
+  args.GetCmdLineArgument("n", num_items);
+  args.GetCmdLineArgument("maxseg", max_segment);
+
+  // Print usage
+  if (args.CheckCmdLineFlag("help"))
+  {
+    printf("%s "
+           "[--n=<input items> "
+           "[--device=<device-id>] "
+           "[--maxseg=<max segment length>] "
+           "[--v] "
+           "\n",
+           argv[0]);
+    exit(0);
+  }
+
+  // Initialize device
+  CubDebugExit(args.DeviceInit());
+
+  // Allocate host arrays
+  int* h_in              = new int[num_items];
+  int* h_reference       = new int[num_items];
+  unsigned char* h_flags = new unsigned char[num_items];
+
+  // Initialize problem and solution
+  Initialize(h_in, h_flags, num_items, max_segment);
+  int num_selected = Solve(h_in, h_flags, h_reference, num_items);
+
+  printf("cub::DevicePartition::Flagged %d items, %d selected (avg distance %d), %d-byte elements\n",
+         num_items,
+         num_selected,
+         (num_selected > 0) ? num_items / num_selected : 0,
+         (int) sizeof(int));
+  fflush(stdout);
+
+  // Allocate problem device arrays
+  int* d_in              = nullptr;
+  unsigned char* d_flags = nullptr;
+
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_in, sizeof(int) * num_items));
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_flags, sizeof(unsigned char) * num_items));
+
+  // Initialize device input
+  CubDebugExit(cudaMemcpy(d_in, h_in, sizeof(int) * num_items, cudaMemcpyHostToDevice));
+  CubDebugExit(cudaMemcpy(d_flags, h_flags, sizeof(unsigned char) * num_items, cudaMemcpyHostToDevice));
+
+  // Allocate device output array and num selected
+  int* d_out              = nullptr;
+  int* d_num_selected_out = nullptr;
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_out, sizeof(int) * num_items));
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_num_selected_out, sizeof(int)));
+
+  // Allocate temporary storage
+  void* d_temp_storage      = nullptr;
+  size_t temp_storage_bytes = 0;
+  CubDebugExit(
+    DevicePartition::Flagged(d_temp_storage, temp_storage_bytes, d_in, d_flags, d_out, d_num_selected_out, num_items));
+  CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));
+
+  // Run
+  CubDebugExit(
+    DevicePartition::Flagged(d_temp_storage, temp_storage_bytes, d_in, d_flags, d_out, d_num_selected_out, num_items));
+
+  // Check for correctness (and display results, if specified)
+  int compare = CompareDeviceResults(h_reference, d_out, num_items, true, g_verbose);
+  printf("\t Data %s ", compare ? "FAIL" : "PASS");
+  compare |= CompareDeviceResults(&num_selected, d_num_selected_out, 1, true, g_verbose);
+  printf("\t Count %s ", compare ? "FAIL" : "PASS");
+  AssertEquals(0, compare);
+
+  // Cleanup
+  if (h_in)
+  {
+    delete[] h_in;
+  }
+  if (h_reference)
+  {
+    delete[] h_reference;
+  }
+  if (d_out)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_out));
+  }
+  if (d_num_selected_out)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_num_selected_out));
+  }
+  if (d_temp_storage)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_temp_storage));
+  }
+  if (d_in)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_in));
+  }
+  if (d_flags)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_flags));
+  }
+
+  printf("\n\n");
+
+  return 0;
+}

cccl_upstream/cub/examples/device/example_device_partition_if.cu

@@ -0,0 +1,234 @@
+// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
+// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+/******************************************************************************
+ * Simple example of DevicePartition::If().
+ *
+ * Partitions items from from a sequence of int keys using a
+ * section functor (greater-than)
+ *
+ * To compile using the command line:
+ *   nvcc -arch=sm_XX example_device_select_if.cu -I../.. -lcudart -O3
+ *
+ ******************************************************************************/
+
+// Ensure printing of CUDA runtime errors to console
+#define CUB_STDERR
+
+#include <cub/device/device_partition.cuh>
+#include <cub/util_allocator.cuh>
+
+#include <cuda/std/limits>
+
+#include <cstdio>
+
+#include "../../test/test_util.h"
+
+using namespace cub;
+
+//---------------------------------------------------------------------
+// Globals, constants and aliases
+//---------------------------------------------------------------------
+
+bool g_verbose = false; // Whether to display input/output to console
+CachingDeviceAllocator g_allocator(true); // Caching allocator for device memory
+
+/// Selection functor type
+struct GreaterThan
+{
+  int compare;
+
+  __host__ __device__ __forceinline__ GreaterThan(int compare)
+      : compare(compare)
+  {}
+
+  __host__ __device__ __forceinline__ bool operator()(const int& a) const
+  {
+    return (a > compare);
+  }
+};
+
+//---------------------------------------------------------------------
+// Test generation
+//---------------------------------------------------------------------
+
+/**
+ * Initialize problem, setting runs of random length chosen from [1..max_segment]
+ */
+void Initialize(int* h_in, int num_items, int max_segment)
+{
+  int key = 0;
+  int i   = 0;
+  while (i < num_items)
+  {
+    // Randomly select number of repeating occurrences uniformly from [1..max_segment]
+    unsigned short bits;
+    RandomBits(bits);
+    const int repeat = cuda::std::max(
+      1,
+      static_cast<int>(static_cast<float>(bits)
+                       * (static_cast<float>(max_segment) / cuda::std::numeric_limits<unsigned short>::max())));
+
+    int j = i;
+    while (j < cuda::std::min(i + repeat, num_items))
+    {
+      h_in[j] = key;
+      j++;
+    }
+
+    i = j;
+    key++;
+  }
+
+  if (g_verbose)
+  {
+    printf("Input:\n");
+    DisplayResults(h_in, num_items);
+    printf("\n\n");
+  }
+}
+
+/**
+ * Solve unique problem
+ */
+template <typename SelectOp>
+int Solve(int* h_in, SelectOp select_op, int* h_reference, int num_items)
+{
+  int num_selected = 0;
+  for (int i = 0; i < num_items; ++i)
+  {
+    if (select_op(h_in[i]))
+    {
+      h_reference[num_selected] = h_in[i];
+      num_selected++;
+    }
+    else
+    {
+      h_reference[num_items - (i - num_selected) - 1] = h_in[i];
+    }
+  }
+
+  return num_selected;
+}
+
+//---------------------------------------------------------------------
+// Main
+//---------------------------------------------------------------------
+
+/**
+ * Main
+ */
+int main(int argc, char** argv)
+{
+  int num_items   = 150;
+  int max_segment = 40; // Maximum segment length
+
+  // Initialize command line
+  CommandLineArgs args(argc, argv);
+  g_verbose = args.CheckCmdLineFlag("v");
+  args.GetCmdLineArgument("n", num_items);
+  args.GetCmdLineArgument("maxseg", max_segment);
+
+  // Print usage
+  if (args.CheckCmdLineFlag("help"))
+  {
+    printf("%s "
+           "[--n=<input items> "
+           "[--device=<device-id>] "
+           "[--maxseg=<max segment length>]"
+           "[--v] "
+           "\n",
+           argv[0]);
+    exit(0);
+  }
+
+  // Initialize device
+  CubDebugExit(args.DeviceInit());
+
+  // Allocate host arrays
+  int* h_in        = new int[num_items];
+  int* h_reference = new int[num_items];
+
+  // DevicePartition a pivot index
+  unsigned int pivot_index;
+  unsigned int max_int = (unsigned int) -1;
+  RandomBits(pivot_index);
+  pivot_index = (unsigned int) ((float(pivot_index) * (float(num_items - 1) / float(max_int))));
+  printf("Pivot idx: %d\n", pivot_index);
+  fflush(stdout);
+
+  // Initialize problem and solution
+  Initialize(h_in, num_items, max_segment);
+  GreaterThan select_op(h_in[pivot_index]);
+
+  int num_selected = Solve(h_in, select_op, h_reference, num_items);
+
+  printf("cub::DevicePartition::If %d items, %d selected (avg run length %d), %d-byte elements\n",
+         num_items,
+         num_selected,
+         (num_selected > 0) ? num_items / num_selected : 0,
+         (int) sizeof(int));
+  fflush(stdout);
+
+  // Allocate problem device arrays
+  int* d_in = nullptr;
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_in, sizeof(int) * num_items));
+
+  // Initialize device input
+  CubDebugExit(cudaMemcpy(d_in, h_in, sizeof(int) * num_items, cudaMemcpyHostToDevice));
+
+  // Allocate device output array and num selected
+  int* d_out              = nullptr;
+  int* d_num_selected_out = nullptr;
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_out, sizeof(int) * num_items));
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_num_selected_out, sizeof(int)));
+
+  // Allocate temporary storage
+  void* d_temp_storage      = nullptr;
+  size_t temp_storage_bytes = 0;
+  CubDebugExit(
+    DevicePartition::If(d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items, select_op));
+  CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));
+
+  // Run
+  CubDebugExit(
+    DevicePartition::If(d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items, select_op));
+
+  // Check for correctness (and display results, if specified)
+  int compare = CompareDeviceResults(h_reference, d_out, num_items, true, g_verbose);
+  printf("\t Data %s ", compare ? "FAIL" : "PASS");
+  compare = compare | CompareDeviceResults(&num_selected, d_num_selected_out, 1, true, g_verbose);
+  printf("\t Count %s ", compare ? "FAIL" : "PASS");
+  AssertEquals(0, compare);
+
+  // Cleanup
+  if (h_in)
+  {
+    delete[] h_in;
+  }
+  if (h_reference)
+  {
+    delete[] h_reference;
+  }
+  if (d_in)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_in));
+  }
+  if (d_out)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_out));
+  }
+  if (d_num_selected_out)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_num_selected_out));
+  }
+  if (d_temp_storage)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_temp_storage));
+  }
+
+  printf("\n\n");
+
+  return 0;
+}

cccl_upstream/cub/examples/device/example_device_radix_sort.cu

@@ -0,0 +1,211 @@
+// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
+// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+/******************************************************************************
+ * Simple example of DeviceRadixSort::SortPairs().
+ *
+ * Sorts an array of float keys paired with a corresponding array of int values.
+ *
+ * To compile using the command line:
+ *   nvcc -arch=sm_XX example_device_radix_sort.cu -I../.. -lcudart -O3
+ *
+ ******************************************************************************/
+
+// Ensure printing of CUDA runtime errors to console
+#define CUB_STDERR
+
+#include <cub/device/device_radix_sort.cuh>
+#include <cub/util_allocator.cuh>
+
+#include <algorithm>
+#include <cstdio>
+
+#include "../../test/test_util.h"
+
+using namespace cub;
+
+//---------------------------------------------------------------------
+// Globals, constants and aliases
+//---------------------------------------------------------------------
+
+bool g_verbose = false; // Whether to display input/output to console
+CachingDeviceAllocator g_allocator(true); // Caching allocator for device memory
+
+//---------------------------------------------------------------------
+// Test generation
+//---------------------------------------------------------------------
+
+/**
+ * Simple key-value pairing for floating point types.
+ * Treats positive and negative zero as equivalent.
+ */
+struct Pair
+{
+  float key;
+  int value;
+
+  bool operator<(const Pair& b) const
+  {
+    return key < b.key;
+  }
+};
+
+/**
+ * Initialize key-value sorting problem.
+ */
+void Initialize(float* h_keys, int* h_values, float* h_reference_keys, int* h_reference_values, int num_items)
+{
+  Pair* h_pairs = new Pair[num_items];
+
+  for (int i = 0; i < num_items; ++i)
+  {
+    RandomBits(h_keys[i]);
+    RandomBits(h_values[i]);
+    h_pairs[i].key   = h_keys[i];
+    h_pairs[i].value = h_values[i];
+  }
+
+  if (g_verbose)
+  {
+    printf("Input keys:\n");
+    DisplayResults(h_keys, num_items);
+    printf("\n\n");
+
+    printf("Input values:\n");
+    DisplayResults(h_values, num_items);
+    printf("\n\n");
+  }
+
+  std::stable_sort(h_pairs, h_pairs + num_items);
+
+  for (int i = 0; i < num_items; ++i)
+  {
+    h_reference_keys[i]   = h_pairs[i].key;
+    h_reference_values[i] = h_pairs[i].value;
+  }
+
+  delete[] h_pairs;
+}
+
+//---------------------------------------------------------------------
+// Main
+//---------------------------------------------------------------------
+
+/**
+ * Main
+ */
+int main(int argc, char** argv)
+{
+  int num_items = 150;
+
+  // Initialize command line
+  CommandLineArgs args(argc, argv);
+  g_verbose = args.CheckCmdLineFlag("v");
+  args.GetCmdLineArgument("n", num_items);
+
+  // Print usage
+  if (args.CheckCmdLineFlag("help"))
+  {
+    printf("%s "
+           "[--n=<input items> "
+           "[--device=<device-id>] "
+           "[--v] "
+           "\n",
+           argv[0]);
+    exit(0);
+  }
+
+  // Initialize device
+  CubDebugExit(args.DeviceInit());
+
+  printf("cub::DeviceRadixSort::SortPairs() %d items (%d-byte keys %d-byte values)\n",
+         num_items,
+         int(sizeof(float)),
+         int(sizeof(int)));
+  fflush(stdout);
+
+  // Allocate host arrays
+  float* h_keys           = new float[num_items];
+  float* h_reference_keys = new float[num_items];
+  int* h_values           = new int[num_items];
+  int* h_reference_values = new int[num_items];
+
+  // Initialize problem and solution on host
+  Initialize(h_keys, h_values, h_reference_keys, h_reference_values, num_items);
+
+  // Allocate device arrays
+  DoubleBuffer<float> d_keys;
+  DoubleBuffer<int> d_values;
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_keys.d_buffers[0], sizeof(float) * num_items));
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_keys.d_buffers[1], sizeof(float) * num_items));
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_values.d_buffers[0], sizeof(int) * num_items));
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_values.d_buffers[1], sizeof(int) * num_items));
+
+  // Allocate temporary storage
+  size_t temp_storage_bytes = 0;
+  void* d_temp_storage      = nullptr;
+
+  CubDebugExit(DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items));
+  CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));
+
+  // Initialize device arrays
+  CubDebugExit(
+    cudaMemcpy(d_keys.d_buffers[d_keys.selector], h_keys, sizeof(float) * num_items, cudaMemcpyHostToDevice));
+  CubDebugExit(
+    cudaMemcpy(d_values.d_buffers[d_values.selector], h_values, sizeof(int) * num_items, cudaMemcpyHostToDevice));
+
+  // Run
+  CubDebugExit(DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items));
+
+  // Check for correctness (and display results, if specified)
+  int compare = CompareDeviceResults(h_reference_keys, d_keys.Current(), num_items, true, g_verbose);
+  printf("\t Compare keys (selector %d): %s\n", d_keys.selector, compare ? "FAIL" : "PASS");
+  AssertEquals(0, compare);
+  compare = CompareDeviceResults(h_reference_values, d_values.Current(), num_items, true, g_verbose);
+  printf("\t Compare values (selector %d): %s\n", d_values.selector, compare ? "FAIL" : "PASS");
+  AssertEquals(0, compare);
+
+  // Cleanup
+  if (h_keys)
+  {
+    delete[] h_keys;
+  }
+  if (h_reference_keys)
+  {
+    delete[] h_reference_keys;
+  }
+  if (h_values)
+  {
+    delete[] h_values;
+  }
+  if (h_reference_values)
+  {
+    delete[] h_reference_values;
+  }
+
+  if (d_keys.d_buffers[0])
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_keys.d_buffers[0]));
+  }
+  if (d_keys.d_buffers[1])
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_keys.d_buffers[1]));
+  }
+  if (d_values.d_buffers[0])
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_values.d_buffers[0]));
+  }
+  if (d_values.d_buffers[1])
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_values.d_buffers[1]));
+  }
+  if (d_temp_storage)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_temp_storage));
+  }
+
+  printf("\n\n");
+
+  return 0;
+}

cccl_upstream/cub/examples/device/example_device_radix_sort_custom.cu

@@ -0,0 +1,233 @@
+// SPDX-FileCopyrightText: Copyright (c) 2023, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+#include <cub/device/device_radix_sort.cuh>
+
+#include <thrust/detail/raw_pointer_cast.h>
+#include <thrust/device_vector.h>
+#include <thrust/host_vector.h>
+
+#include <cuda/std/tuple>
+
+#include <bitset>
+#include <cstdint>
+#include <functional>
+#include <iostream>
+#include <limits>
+#include <type_traits>
+
+#include "cub/block/radix_rank_sort_operations.cuh"
+
+struct custom_t
+{
+  std::uint16_t i;
+  float f;
+};
+
+struct decomposer_t
+{
+  __host__ __device__ //
+    ::cuda::std::tuple<std::uint16_t&, float&>
+    operator()(custom_t& key) const
+  {
+    return {key.i, key.f};
+  }
+};
+
+std::bitset<64> to_binary_representation(custom_t value)
+{
+  std::uint64_t bits{};
+  memcpy(&bits, &value, sizeof(custom_t));
+  return std::bitset<64>{bits};
+}
+
+int main()
+{
+  std::cout << "This example illustrates use of radix sort with custom type.\n";
+  std::cout << "Let's define a simple structure of the following form:\n\n";
+  std::cout << "\tstruct custom_t {\n";
+  std::cout << "\t  std::uint32_t i;\n";
+  std::cout << "\t  float f;\n";
+  std::cout << "\t};\n\n";
+  std::cout << "The `i` field is already stored in the bit-lexicographical order.\n";
+  std::cout << "The `f` field, however, isn't. Therefore, to feed this structure \n";
+  std::cout << "into the radix sort, we have to convert `f` into bit ordered representation.\n";
+  std::cout << "The `custom_t{65535, -4.2f}` has the following binary representation:\n\n";
+
+  auto print_segment = [](std::string msg, std::size_t segment_size, char filler = '-') {
+    std::string spaces((segment_size - msg.size()) / 2 - 1, filler);
+    std::cout << '<' << spaces << msg << spaces << '>';
+  };
+
+  std::cout << '\t';
+  print_segment(" `.f` ", 32);
+  print_segment(" padding -", 16);
+  print_segment(" `.s` ", 16);
+  std::cout << '\n';
+
+  std::cout << "\ts";
+  print_segment(" exp. ", 8);
+  print_segment(" mantissa -", 23);
+  print_segment(" padding -", 16);
+  print_segment(" short -", 16);
+  std::cout << '\n';
+
+  custom_t the_answer{65535, -4.2f};
+  std::cout << '\t' << to_binary_representation(the_answer);
+  std::cout << "\n\t";
+  print_segment(" <----  higher bits  /  lower bits  ----> ", 64, ' ');
+  std::cout << "\n\n";
+
+  std::cout << "Let's say we are trying to compare l={42, -4.2f} with g={42, 4.2f}:\n";
+
+  std::cout << "\n\t";
+  print_segment(" `.f` ", 32);
+  print_segment(" padding -", 16);
+  print_segment(" `.s` ", 16);
+  std::cout << '\n';
+
+  custom_t l{42, -4.2f};
+  custom_t g{42, 4.2f};
+  std::cout << "l:\t" << to_binary_representation(l) << '\n';
+  std::cout << "g:\t" << to_binary_representation(g) << "\n\n";
+
+  std::cout << "As you can see, `l` key happened to be larger in the bit-lexicographical order.\n";
+  std::cout << "Since there's no reflection in C++ (yet), we can't inspect the type and convert \n";
+  std::cout << "each field into the bit-lexicographical order. You can tell CUB how to do that\n";
+  std::cout << "by providing a decomposer for the `custom_t`:\n\n";
+
+  std::cout << "\tstruct decomposer_t \n";
+  std::cout << "\t{\n";
+  std::cout << "\t  __host__ __device__ \n";
+  std::cout << "\t  ::cuda::std::tuple<std::uint16_t&, float&> operator()(custom_t &key) const \n";
+  std::cout << "\t  {\n";
+  std::cout << "\t    return {key.i, key.f};\n";
+  std::cout << "\t  }\n";
+  std::cout << "\t};\n";
+  std::cout << "\n";
+
+  std::cout << "Decomposer allows you to specify which fields are most significant and which\n";
+  std::cout << "are least significant. In our case, `f` is the most significant field and\n";
+  std::cout << "`i` is the least significant field. The decomposer is then used by CUB to convert\n";
+  std::cout << "the `custom_t` into the bit-lexicographical order:\n\n";
+
+  using conversion_policy = cub::detail::radix::traits_t<custom_t>::bit_ordered_conversion_policy;
+  l                       = conversion_policy::to_bit_ordered(decomposer_t{}, l);
+  g                       = conversion_policy::to_bit_ordered(decomposer_t{}, g);
+
+  std::cout << "\n\t";
+  print_segment(" `.f` ", 32);
+  print_segment(" padding -", 16);
+  print_segment(" `.s` ", 16);
+  std::cout << '\n';
+
+  std::cout << "l:\t" << to_binary_representation(l) << '\n';
+  std::cout << "g:\t" << to_binary_representation(g) << "\n\n";
+
+  std::cout << '\n';
+  std::cout << "As you can see, `g` is now actually larger than `l` in the bit-lexicographical order.\n";
+  std::cout << "After binning, CUB is able to restore the original key:\n\n";
+
+  l = conversion_policy::from_bit_ordered(decomposer_t{}, l);
+  g = conversion_policy::from_bit_ordered(decomposer_t{}, g);
+
+  std::cout << "\n\t";
+  print_segment(" `.f` ", 32);
+  print_segment(" padding -", 16);
+  print_segment(" `.s` ", 16);
+  std::cout << '\n';
+
+  std::cout << "l:\t" << to_binary_representation(l) << '\n';
+  std::cout << "g:\t" << to_binary_representation(g) << "\n\n";
+
+  using inversion_policy = cub::detail::radix::traits_t<custom_t>::bit_ordered_inversion_policy;
+  std::cout << '\n';
+  std::cout << "We are also able to inverse differentiating bits:\n";
+
+  l = inversion_policy::inverse(decomposer_t{}, l);
+  g = inversion_policy::inverse(decomposer_t{}, g);
+
+  std::cout << "\n\t";
+  print_segment(" `.f` ", 32);
+  print_segment(" padding -", 16);
+  print_segment(" `.s` ", 16);
+  std::cout << '\n';
+
+  std::cout << "l:\t" << to_binary_representation(l) << '\n';
+  std::cout << "g:\t" << to_binary_representation(g) << "\n\n";
+
+  std::cout << '\n';
+  std::cout << "We as well can compute the minimal and minimal / maximal keys:\n";
+
+  l = cub::detail::radix::traits_t<custom_t>::min_raw_binary_key(decomposer_t{});
+  g = cub::detail::radix::traits_t<custom_t>::max_raw_binary_key(decomposer_t{});
+
+  std::cout << "\n\t";
+  print_segment(" `.f` ", 32);
+  print_segment(" padding -", 16);
+  print_segment(" `.s` ", 16);
+  std::cout << '\n';
+
+  std::cout << "l:\t" << to_binary_representation(l) << '\n';
+  std::cout << "g:\t" << to_binary_representation(g) << "\n\n";
+
+  std::cout << "We can even compute the number of differentiating bits:\n\n";
+
+  std::cout << "end:\t";
+  std::cout << cub::detail::radix::traits_t<custom_t>::default_end_bit(decomposer_t{});
+  std::cout << '\n';
+  std::cout << "size:\t";
+  std::cout << sizeof(custom_t) * CHAR_BIT;
+  std::cout << "\n\n";
+
+  std::cout << "All of these operations are used behind the scenes by CUB to sort custom types:\n\n";
+
+  constexpr int num_items            = 6;
+  thrust::device_vector<custom_t> in = {{4, +2.5f}, {0, -2.5f}, {3, +1.1f}, {1, +0.0f}, {2, -0.0f}, {5, +3.7f}};
+
+  std::cout << "in:\n";
+  for (custom_t key : in)
+  {
+    std::cout << "\t{.i = " << key.i << ", .f = " << key.f << "},\n";
+  }
+
+  thrust::device_vector<custom_t> out(num_items);
+
+  const custom_t* d_in = thrust::raw_pointer_cast(in.data());
+  custom_t* d_out      = thrust::raw_pointer_cast(out.data());
+
+  // 1) Get temp storage size
+  std::uint8_t* d_temp_storage{};
+  std::size_t temp_storage_bytes{};
+
+  cub::DeviceRadixSort::SortKeys(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items, decomposer_t{});
+
+  // 2) Allocate temp storage
+  thrust::device_vector<std::uint8_t> temp_storage(temp_storage_bytes);
+  d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());
+
+  // 3) Sort keys
+  cub::DeviceRadixSort::SortKeys(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items, decomposer_t{});
+  cudaDeviceSynchronize();
+
+  std::cout << "\n";
+  std::cout << "sort:\n";
+  std::cout << "\n";
+
+  std::cout << "\tcub::DeviceRadixSort::SortKeys(d_temp_storage,\n";
+  std::cout << "\t                               temp_storage_bytes,\n";
+  std::cout << "\t                               d_in,\n";
+  std::cout << "\t                               d_out,\n";
+  std::cout << "\t                               num_items,\n";
+  std::cout << "\t                               decomposer_t{});\n\n";
+
+  std::cout << "out:\n";
+  for (custom_t key : out)
+  {
+    std::cout << "\t{.i = " << key.i << ", .f = " << key.f << "},\n";
+  }
+
+  std::cout << '\n';
+  std::cout << "If you have any issues with radix sort support of custom types, \n";
+  std::cout << "please feel free to use this example to identify the problem.\n\n";
+}

cccl_upstream/cub/examples/device/example_device_reduce.cu

@@ -0,0 +1,164 @@
+// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
+// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+/******************************************************************************
+ * Simple example of DeviceReduce::Sum().
+ *
+ * Sums an array of int keys.
+ *
+ * To compile using the command line:
+ *   nvcc -arch=sm_XX example_device_reduce.cu -I../.. -lcudart -O3
+ *
+ ******************************************************************************/
+
+// Ensure printing of CUDA runtime errors to console
+#define CUB_STDERR
+
+#include <cub/device/device_reduce.cuh>
+#include <cub/util_allocator.cuh>
+
+#include <cstdio>
+
+#include "../../test/test_util.h"
+
+using namespace cub;
+
+//---------------------------------------------------------------------
+// Globals, constants and aliases
+//---------------------------------------------------------------------
+
+bool g_verbose = false; // Whether to display input/output to console
+CachingDeviceAllocator g_allocator(true); // Caching allocator for device memory
+
+//---------------------------------------------------------------------
+// Test generation
+//---------------------------------------------------------------------
+
+/**
+ * Initialize problem
+ */
+void Initialize(int* h_in, int num_items)
+{
+  for (int i = 0; i < num_items; ++i)
+  {
+    h_in[i] = i;
+  }
+
+  if (g_verbose)
+  {
+    printf("Input:\n");
+    DisplayResults(h_in, num_items);
+    printf("\n\n");
+  }
+}
+
+/**
+ * Compute solution
+ */
+void Solve(int* h_in, int& h_reference, int num_items)
+{
+  for (int i = 0; i < num_items; ++i)
+  {
+    if (i == 0)
+    {
+      h_reference = h_in[0];
+    }
+    else
+    {
+      h_reference += h_in[i];
+    }
+  }
+}
+
+//---------------------------------------------------------------------
+// Main
+//---------------------------------------------------------------------
+
+/**
+ * Main
+ */
+int main(int argc, char** argv)
+{
+  int num_items = 150;
+
+  // Initialize command line
+  CommandLineArgs args(argc, argv);
+  g_verbose = args.CheckCmdLineFlag("v");
+  args.GetCmdLineArgument("n", num_items);
+
+  // Print usage
+  if (args.CheckCmdLineFlag("help"))
+  {
+    printf("%s "
+           "[--n=<input items> "
+           "[--device=<device-id>] "
+           "[--v] "
+           "\n",
+           argv[0]);
+    exit(0);
+  }
+
+  // Initialize device
+  CubDebugExit(args.DeviceInit());
+
+  printf("cub::DeviceReduce::Sum() %d items (%d-byte elements)\n", num_items, (int) sizeof(int));
+  fflush(stdout);
+
+  // Allocate host arrays
+  int* h_in = new int[num_items];
+  int h_reference{};
+
+  // Initialize problem and solution
+  Initialize(h_in, num_items);
+  Solve(h_in, h_reference, num_items);
+
+  // Allocate problem device arrays
+  int* d_in = nullptr;
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_in, sizeof(int) * num_items));
+
+  // Initialize device input
+  CubDebugExit(cudaMemcpy(d_in, h_in, sizeof(int) * num_items, cudaMemcpyHostToDevice));
+
+  // Allocate device output array
+  int* d_out = nullptr;
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_out, sizeof(int) * 1));
+
+  // example-begin temp-storage-query
+  // Request and allocate temporary storage
+  void* d_temp_storage      = nullptr;
+  size_t temp_storage_bytes = 0;
+  CubDebugExit(DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items));
+  CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));
+
+  // Run
+  CubDebugExit(DeviceReduce::Sum(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items));
+  // example-end temp-storage-query
+
+  // Check for correctness (and display results, if specified)
+  int compare = CompareDeviceResults(&h_reference, d_out, 1, g_verbose, g_verbose);
+  printf("\t%s", compare ? "FAIL" : "PASS");
+  AssertEquals(0, compare);
+
+  // Cleanup
+  if (h_in)
+  {
+    delete[] h_in;
+  }
+  if (d_in)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_in));
+  }
+  if (d_out)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_out));
+  }
+  if (d_temp_storage)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_temp_storage));
+  }
+
+  printf("\n\n");
+
+  return 0;
+}

cccl_upstream/cub/examples/device/example_device_reduce_env.cu

@@ -0,0 +1,120 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+/// Simple example of DeviceReduce::Sum() using an environment
+/// Sums an array of int keys.
+
+// Ensure printing of CUDA runtime errors to console
+#define CUB_STDERR
+
+#include <cub/device/device_reduce.cuh>
+
+#include <thrust/device_vector.h>
+
+#include <cuda/devices>
+#include <cuda/memory_pool>
+#include <cuda/stream>
+
+#include <cstdio>
+
+#include "../../test/test_util.h"
+
+bool g_verbose = false; // Whether to display input/output to console
+
+void Initialize(int* h_in, int num_items)
+{
+  for (int i = 0; i < num_items; ++i)
+  {
+    h_in[i] = i;
+  }
+
+  if (g_verbose)
+  {
+    printf("Input:\n");
+    DisplayResults(h_in, num_items);
+    printf("\n\n");
+  }
+}
+
+void Solve(int* h_in, int& h_reference, int num_items)
+{
+  for (int i = 0; i < num_items; ++i)
+  {
+    if (i == 0)
+    {
+      h_reference = h_in[0];
+    }
+    else
+    {
+      h_reference += h_in[i];
+    }
+  }
+}
+
+int main(int argc, char** argv)
+{
+  // Initialize command line and print usage
+  CommandLineArgs args(argc, argv);
+  if (args.CheckCmdLineFlag("help"))
+  {
+    printf("%s "
+           "[--n=<input items> "
+           "[--device=<device-id>] "
+           "[--v] "
+           "\n",
+           argv[0]);
+    std::exit(0);
+  }
+
+  // Parse command line options
+  int num_items      = 150;
+  int device_ordinal = 0;
+  g_verbose          = args.CheckCmdLineFlag("v");
+  args.GetCmdLineArgument("n", num_items);
+  args.GetCmdLineArgument("d", device_ordinal);
+
+  // example-begin env-overload-setup
+  // 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};
+  // example-end env-overload-setup
+
+  printf("cub::DeviceReduce::Sum() %d items (%d-byte elements)\n", num_items, (int) sizeof(int));
+  fflush(stdout);
+
+  // Allocate host arrays
+  std::vector<int> h_in(num_items);
+  int h_reference = 0;
+
+  // Initialize problem and solution
+  Initialize(h_in.data(), num_items);
+  Solve(h_in.data(), h_reference, num_items);
+
+  // Allocate problem device arrays
+  auto d_in = thrust::device_vector<int>(num_items, thrust::no_init);
+
+  // Initialize device input
+  thrust::copy(h_in.begin(), h_in.end(), d_in.begin());
+
+  // Allocate device output array
+  auto d_out = thrust::device_vector<int>(1);
+
+  // example-begin env-overload-run
+  // Run
+  CubDebugExit(cub::DeviceReduce::Sum(d_in.data(), d_out.data(), num_items, env));
+  // example-end env-overload-run
+
+  // Check for correctness
+  // Check for correctness (and display results, if specified)
+  const int compare =
+    CompareDeviceResults(&h_reference, thrust::raw_pointer_cast(d_out.data()), 1, g_verbose, g_verbose);
+  printf("\t%s", compare ? "FAIL" : "PASS");
+  printf("\n\n");
+
+  return 0;
+}

cccl_upstream/cub/examples/device/example_device_scan.cu

@@ -0,0 +1,166 @@
+// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
+// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+/******************************************************************************
+ * Simple example of DeviceScan::ExclusiveSum().
+ *
+ * Computes an exclusive sum of int keys.
+ *
+ * To compile using the command line:
+ *   nvcc -arch=sm_XX example_device_scan.cu -I../.. -lcudart -O3
+ *
+ ******************************************************************************/
+
+// Ensure printing of CUDA runtime errors to console
+#define CUB_STDERR
+
+#include <cub/device/device_scan.cuh>
+#include <cub/util_allocator.cuh>
+
+#include <cstdio>
+
+#include "../../test/test_util.h"
+
+using namespace cub;
+
+//---------------------------------------------------------------------
+// Globals, constants and aliases
+//---------------------------------------------------------------------
+
+bool g_verbose = false; // Whether to display input/output to console
+CachingDeviceAllocator g_allocator(true); // Caching allocator for device memory
+
+//---------------------------------------------------------------------
+// Test generation
+//---------------------------------------------------------------------
+
+/**
+ * Initialize problem
+ */
+void Initialize(int* h_in, int num_items)
+{
+  for (int i = 0; i < num_items; ++i)
+  {
+    h_in[i] = i;
+  }
+
+  if (g_verbose)
+  {
+    printf("Input:\n");
+    DisplayResults(h_in, num_items);
+    printf("\n\n");
+  }
+}
+
+/**
+ * Solve exclusive-scan problem
+ */
+int Solve(int* h_in, int* h_reference, int num_items)
+{
+  int inclusive = 0;
+  int aggregate = 0;
+
+  for (int i = 0; i < num_items; ++i)
+  {
+    h_reference[i] = inclusive;
+    inclusive += h_in[i];
+    aggregate += h_in[i];
+  }
+
+  return aggregate;
+}
+
+//---------------------------------------------------------------------
+// Main
+//---------------------------------------------------------------------
+
+/**
+ * Main
+ */
+int main(int argc, char** argv)
+{
+  int num_items = 150;
+
+  // Initialize command line
+  CommandLineArgs args(argc, argv);
+  g_verbose = args.CheckCmdLineFlag("v");
+  args.GetCmdLineArgument("n", num_items);
+
+  // Print usage
+  if (args.CheckCmdLineFlag("help"))
+  {
+    printf("%s "
+           "[--n=<input items> "
+           "[--device=<device-id>] "
+           "[--v] "
+           "\n",
+           argv[0]);
+    exit(0);
+  }
+
+  // Initialize device
+  CubDebugExit(args.DeviceInit());
+
+  printf("cub::DeviceScan::ExclusiveSum %d items (%d-byte elements)\n", num_items, (int) sizeof(int));
+  fflush(stdout);
+
+  // Allocate host arrays
+  int* h_in        = new int[num_items];
+  int* h_reference = new int[num_items];
+
+  // Initialize problem and solution
+  Initialize(h_in, num_items);
+  Solve(h_in, h_reference, num_items);
+
+  // Allocate problem device arrays
+  int* d_in = nullptr;
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_in, sizeof(int) * num_items));
+
+  // Initialize device input
+  CubDebugExit(cudaMemcpy(d_in, h_in, sizeof(int) * num_items, cudaMemcpyHostToDevice));
+
+  // Allocate device output array
+  int* d_out = nullptr;
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_out, sizeof(int) * num_items));
+
+  // Allocate temporary storage
+  void* d_temp_storage      = nullptr;
+  size_t temp_storage_bytes = 0;
+  CubDebugExit(DeviceScan::ExclusiveSum(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items));
+  CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));
+
+  // Run
+  CubDebugExit(DeviceScan::ExclusiveSum(d_temp_storage, temp_storage_bytes, d_in, d_out, num_items));
+
+  // Check for correctness (and display results, if specified)
+  int compare = CompareDeviceResults(h_reference, d_out, num_items, true, g_verbose);
+  printf("\t%s", compare ? "FAIL" : "PASS");
+  AssertEquals(0, compare);
+
+  // Cleanup
+  if (h_in)
+  {
+    delete[] h_in;
+  }
+  if (h_reference)
+  {
+    delete[] h_reference;
+  }
+  if (d_in)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_in));
+  }
+  if (d_out)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_out));
+  }
+  if (d_temp_storage)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_temp_storage));
+  }
+
+  printf("\n\n");
+
+  return 0;
+}

cccl_upstream/cub/examples/device/example_device_select_flagged.cu

@@ -0,0 +1,222 @@
+// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
+// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+/******************************************************************************
+ * Simple example of DeviceSelect::Flagged().
+ *
+ * Selects flagged items from from a sequence of int keys using a
+ * corresponding sequence of unsigned char flags.
+ *
+ * To compile using the command line:
+ *   nvcc -arch=sm_XX example_device_select_flagged.cu -I../.. -lcudart -O3
+ *
+ ******************************************************************************/
+
+// Ensure printing of CUDA runtime errors to console
+#define CUB_STDERR
+
+#include <cub/device/device_select.cuh>
+#include <cub/util_allocator.cuh>
+
+#include <cuda/std/limits>
+
+#include <cstdio>
+
+#include "../../test/test_util.h"
+
+using namespace cub;
+
+//---------------------------------------------------------------------
+// Globals, constants and aliases
+//---------------------------------------------------------------------
+
+bool g_verbose = false; // Whether to display input/output to console
+CachingDeviceAllocator g_allocator(true); // Caching allocator for device memory
+
+//---------------------------------------------------------------------
+// Test generation
+//---------------------------------------------------------------------
+
+/**
+ * Initialize problem, setting flags at distances of random length
+ * chosen from [1..max_segment]
+ */
+void Initialize(int* h_in, unsigned char* h_flags, int num_items, int max_segment)
+{
+  int key = 0;
+  int i   = 0;
+  while (i < num_items)
+  {
+    // Select number of repeating occurrences
+    unsigned short bits;
+    RandomBits(bits);
+    const int repeat = cuda::std::max(
+      1,
+      static_cast<int>(static_cast<float>(bits)
+                       * (static_cast<float>(max_segment) / cuda::std::numeric_limits<unsigned short>::max())));
+
+    int j = i;
+    while (j < cuda::std::min(i + repeat, num_items))
+    {
+      h_flags[j] = 0;
+      h_in[j]    = key;
+      j++;
+    }
+
+    h_flags[i] = 1;
+    i          = j;
+    key++;
+  }
+
+  if (g_verbose)
+  {
+    printf("Input:\n");
+    DisplayResults(h_in, num_items);
+    printf("Flags:\n");
+    DisplayResults(h_flags, num_items);
+    printf("\n\n");
+  }
+}
+
+/**
+ * Solve unique problem
+ */
+int Solve(int* h_in, unsigned char* h_flags, int* h_reference, int num_items)
+{
+  int num_selected = 0;
+  for (int i = 0; i < num_items; ++i)
+  {
+    if (h_flags[i])
+    {
+      h_reference[num_selected] = h_in[i];
+      num_selected++;
+    }
+    else
+    {
+      h_reference[num_items - (i - num_selected) - 1] = h_in[i];
+    }
+  }
+
+  return num_selected;
+}
+
+//---------------------------------------------------------------------
+// Main
+//---------------------------------------------------------------------
+
+/**
+ * Main
+ */
+int main(int argc, char** argv)
+{
+  int num_items   = 150;
+  int max_segment = 40; // Maximum segment length
+
+  // Initialize command line
+  CommandLineArgs args(argc, argv);
+  g_verbose = args.CheckCmdLineFlag("v");
+  args.GetCmdLineArgument("n", num_items);
+  args.GetCmdLineArgument("maxseg", max_segment);
+
+  // Print usage
+  if (args.CheckCmdLineFlag("help"))
+  {
+    printf("%s "
+           "[--n=<input items> "
+           "[--device=<device-id>] "
+           "[--maxseg=<max segment length>] "
+           "[--v] "
+           "\n",
+           argv[0]);
+    exit(0);
+  }
+
+  // Initialize device
+  CubDebugExit(args.DeviceInit());
+
+  // Allocate host arrays
+  int* h_in              = new int[num_items];
+  int* h_reference       = new int[num_items];
+  unsigned char* h_flags = new unsigned char[num_items];
+
+  // Initialize problem and solution
+  Initialize(h_in, h_flags, num_items, max_segment);
+  int num_selected = Solve(h_in, h_flags, h_reference, num_items);
+
+  printf("cub::DeviceSelect::Flagged %d items, %d selected (avg distance %d), %d-byte elements\n",
+         num_items,
+         num_selected,
+         (num_selected > 0) ? num_items / num_selected : 0,
+         (int) sizeof(int));
+  fflush(stdout);
+
+  // Allocate problem device arrays
+  int* d_in              = nullptr;
+  unsigned char* d_flags = nullptr;
+
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_in, sizeof(int) * num_items));
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_flags, sizeof(unsigned char) * num_items));
+
+  // Initialize device input
+  CubDebugExit(cudaMemcpy(d_in, h_in, sizeof(int) * num_items, cudaMemcpyHostToDevice));
+  CubDebugExit(cudaMemcpy(d_flags, h_flags, sizeof(unsigned char) * num_items, cudaMemcpyHostToDevice));
+
+  // Allocate device output array and num selected
+  int* d_out              = nullptr;
+  int* d_num_selected_out = nullptr;
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_out, sizeof(int) * num_items));
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_num_selected_out, sizeof(int)));
+
+  // Allocate temporary storage
+  void* d_temp_storage      = nullptr;
+  size_t temp_storage_bytes = 0;
+  CubDebugExit(
+    DeviceSelect::Flagged(d_temp_storage, temp_storage_bytes, d_in, d_flags, d_out, d_num_selected_out, num_items));
+  CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));
+
+  // Run
+  CubDebugExit(
+    DeviceSelect::Flagged(d_temp_storage, temp_storage_bytes, d_in, d_flags, d_out, d_num_selected_out, num_items));
+
+  // Check for correctness (and display results, if specified)
+  int compare = CompareDeviceResults(h_reference, d_out, num_selected, true, g_verbose);
+  printf("\t Data %s ", compare ? "FAIL" : "PASS");
+  compare |= CompareDeviceResults(&num_selected, d_num_selected_out, 1, true, g_verbose);
+  printf("\t Count %s ", compare ? "FAIL" : "PASS");
+  AssertEquals(0, compare);
+
+  // Cleanup
+  if (h_in)
+  {
+    delete[] h_in;
+  }
+  if (h_reference)
+  {
+    delete[] h_reference;
+  }
+  if (d_out)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_out));
+  }
+  if (d_num_selected_out)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_num_selected_out));
+  }
+  if (d_temp_storage)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_temp_storage));
+  }
+  if (d_in)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_in));
+  }
+  if (d_flags)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_flags));
+  }
+
+  printf("\n\n");
+
+  return 0;
+}

cccl_upstream/cub/examples/device/example_device_select_if.cu

@@ -0,0 +1,234 @@
+// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
+// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+/******************************************************************************
+ * Simple example of DeviceSelect::If().
+ *
+ * Selects items from from a sequence of int keys using a
+ * section functor (greater-than)
+ *
+ * To compile using the command line:
+ *   nvcc -arch=sm_XX example_device_select_if.cu -I../.. -lcudart -O3
+ *
+ ******************************************************************************/
+
+// Ensure printing of CUDA runtime errors to console
+#define CUB_STDERR
+
+#include <cub/device/device_select.cuh>
+#include <cub/util_allocator.cuh>
+
+#include <cuda/std/limits>
+
+#include <cstdio>
+
+#include "../../test/test_util.h"
+
+using namespace cub;
+
+//---------------------------------------------------------------------
+// Globals, constants and aliases
+//---------------------------------------------------------------------
+
+bool g_verbose = false; // Whether to display input/output to console
+CachingDeviceAllocator g_allocator(true); // Caching allocator for device memory
+
+/// Selection functor type
+struct GreaterThan
+{
+  int compare;
+
+  __host__ __device__ __forceinline__ GreaterThan(int compare)
+      : compare(compare)
+  {}
+
+  __host__ __device__ __forceinline__ bool operator()(const int& a) const
+  {
+    return (a > compare);
+  }
+};
+
+//---------------------------------------------------------------------
+// Test generation
+//---------------------------------------------------------------------
+
+/**
+ * Initialize problem, setting runs of random length chosen from [1..max_segment]
+ */
+void Initialize(int* h_in, int num_items, int max_segment)
+{
+  int key = 0;
+  int i   = 0;
+  while (i < num_items)
+  {
+    // Randomly select number of repeating occurrences uniformly from [1..max_segment]
+    unsigned short bits;
+    RandomBits(bits);
+    const int repeat = cuda::std::max(
+      1,
+      static_cast<int>(static_cast<float>(bits)
+                       * (static_cast<float>(max_segment) / cuda::std::numeric_limits<unsigned short>::max())));
+
+    int j = i;
+    while (j < cuda::std::min(i + repeat, num_items))
+    {
+      h_in[j] = key;
+      j++;
+    }
+
+    i = j;
+    key++;
+  }
+
+  if (g_verbose)
+  {
+    printf("Input:\n");
+    DisplayResults(h_in, num_items);
+    printf("\n\n");
+  }
+}
+
+/**
+ * Solve unique problem
+ */
+template <typename SelectOp>
+int Solve(int* h_in, SelectOp select_op, int* h_reference, int num_items)
+{
+  int num_selected = 0;
+  for (int i = 0; i < num_items; ++i)
+  {
+    if (select_op(h_in[i]))
+    {
+      h_reference[num_selected] = h_in[i];
+      num_selected++;
+    }
+    else
+    {
+      h_reference[num_items - (i - num_selected) - 1] = h_in[i];
+    }
+  }
+
+  return num_selected;
+}
+
+//---------------------------------------------------------------------
+// Main
+//---------------------------------------------------------------------
+
+/**
+ * Main
+ */
+int main(int argc, char** argv)
+{
+  int num_items   = 150;
+  int max_segment = 40; // Maximum segment length
+
+  // Initialize command line
+  CommandLineArgs args(argc, argv);
+  g_verbose = args.CheckCmdLineFlag("v");
+  args.GetCmdLineArgument("n", num_items);
+  args.GetCmdLineArgument("maxseg", max_segment);
+
+  // Print usage
+  if (args.CheckCmdLineFlag("help"))
+  {
+    printf("%s "
+           "[--n=<input items> "
+           "[--device=<device-id>] "
+           "[--maxseg=<max segment length>]"
+           "[--v] "
+           "\n",
+           argv[0]);
+    exit(0);
+  }
+
+  // Initialize device
+  CubDebugExit(args.DeviceInit());
+
+  // Allocate host arrays
+  int* h_in        = new int[num_items];
+  int* h_reference = new int[num_items];
+
+  // Select a pivot index
+  unsigned int pivot_index;
+  unsigned int max_int = (unsigned int) -1;
+  RandomBits(pivot_index);
+  pivot_index = (unsigned int) ((float(pivot_index) * (float(num_items - 1) / float(max_int))));
+  printf("Pivot idx: %d\n", pivot_index);
+  fflush(stdout);
+
+  // Initialize problem and solution
+  Initialize(h_in, num_items, max_segment);
+  GreaterThan select_op(h_in[pivot_index]);
+
+  int num_selected = Solve(h_in, select_op, h_reference, num_items);
+
+  printf("cub::DeviceSelect::If %d items, %d selected (avg run length %d), %d-byte elements\n",
+         num_items,
+         num_selected,
+         (num_selected > 0) ? num_items / num_selected : 0,
+         (int) sizeof(int));
+  fflush(stdout);
+
+  // Allocate problem device arrays
+  int* d_in = nullptr;
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_in, sizeof(int) * num_items));
+
+  // Initialize device input
+  CubDebugExit(cudaMemcpy(d_in, h_in, sizeof(int) * num_items, cudaMemcpyHostToDevice));
+
+  // Allocate device output array and num selected
+  int* d_out              = nullptr;
+  int* d_num_selected_out = nullptr;
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_out, sizeof(int) * num_items));
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_num_selected_out, sizeof(int)));
+
+  // Allocate temporary storage
+  void* d_temp_storage      = nullptr;
+  size_t temp_storage_bytes = 0;
+  CubDebugExit(
+    DeviceSelect::If(d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items, select_op));
+  CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));
+
+  // Run
+  CubDebugExit(
+    DeviceSelect::If(d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items, select_op));
+
+  // Check for correctness (and display results, if specified)
+  int compare = CompareDeviceResults(h_reference, d_out, num_selected, true, g_verbose);
+  printf("\t Data %s ", compare ? "FAIL" : "PASS");
+  compare = compare | CompareDeviceResults(&num_selected, d_num_selected_out, 1, true, g_verbose);
+  printf("\t Count %s ", compare ? "FAIL" : "PASS");
+  AssertEquals(0, compare);
+
+  // Cleanup
+  if (h_in)
+  {
+    delete[] h_in;
+  }
+  if (h_reference)
+  {
+    delete[] h_reference;
+  }
+  if (d_in)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_in));
+  }
+  if (d_out)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_out));
+  }
+  if (d_num_selected_out)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_num_selected_out));
+  }
+  if (d_temp_storage)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_temp_storage));
+  }
+
+  printf("\n\n");
+
+  return 0;
+}

cccl_upstream/cub/examples/device/example_device_select_unique.cu

@@ -0,0 +1,209 @@
+// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
+// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+/******************************************************************************
+ * Simple example of DeviceSelect::Unique().
+ *
+ * Selects the first element from each run of identical values from a sequence
+ * of int keys.
+ *
+ * To compile using the command line:
+ *   nvcc -arch=sm_XX example_device_select_unique.cu -I../.. -lcudart -O3
+ *
+ ******************************************************************************/
+
+// Ensure printing of CUDA runtime errors to console
+#define CUB_STDERR
+
+#include <cub/device/device_select.cuh>
+#include <cub/util_allocator.cuh>
+
+#include <cuda/std/limits>
+
+#include <cstdio>
+
+#include "../../test/test_util.h"
+
+using namespace cub;
+
+//---------------------------------------------------------------------
+// Globals, constants and aliases
+//---------------------------------------------------------------------
+
+bool g_verbose = false; // Whether to display input/output to console
+CachingDeviceAllocator g_allocator(true); // Caching allocator for device memory
+
+//---------------------------------------------------------------------
+// Test generation
+//---------------------------------------------------------------------
+
+/**
+ * Initialize problem, setting runs of random length chosen from [1..max_segment]
+ */
+void Initialize(int* h_in, int num_items, int max_segment)
+{
+  int key = 0;
+  int i   = 0;
+  while (i < num_items)
+  {
+    // Randomly select number of repeating occurrences uniformly from [1..max_segment]
+    unsigned short bits;
+    RandomBits(bits);
+    const int repeat = cuda::std::max(
+      1,
+      static_cast<int>(static_cast<float>(bits)
+                       * (static_cast<float>(max_segment) / cuda::std::numeric_limits<unsigned short>::max())));
+    ;
+
+    int j = i;
+    while (j < cuda::std::min(i + repeat, num_items))
+    {
+      h_in[j] = key;
+      j++;
+    }
+
+    i = j;
+    key++;
+  }
+
+  if (g_verbose)
+  {
+    printf("Input:\n");
+    DisplayResults(h_in, num_items);
+    printf("\n\n");
+  }
+}
+
+/**
+ * Solve unique problem
+ */
+int Solve(int* h_in, int* h_reference, int num_items)
+{
+  int num_selected = 0;
+  if (num_items > 0)
+  {
+    h_reference[num_selected] = h_in[0];
+    num_selected++;
+  }
+
+  for (int i = 1; i < num_items; ++i)
+  {
+    if (h_in[i] != h_in[i - 1])
+    {
+      h_reference[num_selected] = h_in[i];
+      num_selected++;
+    }
+  }
+
+  return num_selected;
+}
+
+//---------------------------------------------------------------------
+// Main
+//---------------------------------------------------------------------
+
+/**
+ * Main
+ */
+int main(int argc, char** argv)
+{
+  int num_items   = 150;
+  int max_segment = 40; // Maximum segment length
+
+  // Initialize command line
+  CommandLineArgs args(argc, argv);
+  g_verbose = args.CheckCmdLineFlag("v");
+  args.GetCmdLineArgument("n", num_items);
+  args.GetCmdLineArgument("maxseg", max_segment);
+
+  // Print usage
+  if (args.CheckCmdLineFlag("help"))
+  {
+    printf("%s "
+           "[--n=<input items> "
+           "[--device=<device-id>] "
+           "[--maxseg=<max segment length>]"
+           "[--v] "
+           "\n",
+           argv[0]);
+    exit(0);
+  }
+
+  // Initialize device
+  CubDebugExit(args.DeviceInit());
+
+  // Allocate host arrays
+  int* h_in        = new int[num_items];
+  int* h_reference = new int[num_items];
+
+  // Initialize problem and solution
+  Initialize(h_in, num_items, max_segment);
+  int num_selected = Solve(h_in, h_reference, num_items);
+
+  printf("cub::DeviceSelect::Unique %d items (%d-byte elements), %d selected (avg run length %d)\n",
+         num_items,
+         (int) sizeof(int),
+         num_selected,
+         num_items / num_selected);
+  fflush(stdout);
+
+  // Allocate problem device arrays
+  int* d_in = nullptr;
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_in, sizeof(int) * num_items));
+
+  // Initialize device input
+  CubDebugExit(cudaMemcpy(d_in, h_in, sizeof(int) * num_items, cudaMemcpyHostToDevice));
+
+  // Allocate device output array and num selected
+  int* d_out              = nullptr;
+  int* d_num_selected_out = nullptr;
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_out, sizeof(int) * num_items));
+  CubDebugExit(g_allocator.DeviceAllocate((void**) &d_num_selected_out, sizeof(int)));
+
+  // Allocate temporary storage
+  void* d_temp_storage      = nullptr;
+  size_t temp_storage_bytes = 0;
+  CubDebugExit(DeviceSelect::Unique(d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items));
+  CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));
+
+  // Run
+  CubDebugExit(DeviceSelect::Unique(d_temp_storage, temp_storage_bytes, d_in, d_out, d_num_selected_out, num_items));
+
+  // Check for correctness (and display results, if specified)
+  int compare = CompareDeviceResults(h_reference, d_out, num_selected, true, g_verbose);
+  printf("\t Data %s ", compare ? "FAIL" : "PASS");
+  compare = compare | CompareDeviceResults(&num_selected, d_num_selected_out, 1, true, g_verbose);
+  printf("\t Count %s ", compare ? "FAIL" : "PASS");
+  AssertEquals(0, compare);
+
+  // Cleanup
+  if (h_in)
+  {
+    delete[] h_in;
+  }
+  if (h_reference)
+  {
+    delete[] h_reference;
+  }
+  if (d_in)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_in));
+  }
+  if (d_out)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_out));
+  }
+  if (d_num_selected_out)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_num_selected_out));
+  }
+  if (d_temp_storage)
+  {
+    CubDebugExit(g_allocator.DeviceFree(d_temp_storage));
+  }
+
+  printf("\n\n");
+
+  return 0;
+}

cccl_upstream/cub/examples/device/example_device_sort_find_non_trivial_runs.cu

@@ -0,0 +1,384 @@
+// SPDX-FileCopyrightText: Copyright (c) 2011, Duane Merrill. All rights reserved.
+// SPDX-FileCopyrightText: Copyright (c) 2011-2018, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+/******************************************************************************
+ * Simple example of sorting a sequence of keys and values (each pair is a
+ * randomly-selected int32 paired with its original offset in the unsorted sequence), and then
+ * isolating all maximal, non-trivial (having length > 1) "runs" of duplicates.
+ *
+ * To compile using the command line:
+ *   nvcc -arch=sm_XX example_device_sort_find_non_trivial_runs.cu -I../.. -lcudart -O3
+ *
+ ******************************************************************************/
+
+// Ensure printing of CUDA runtime errors to console
+#define CUB_STDERR
+
+#include <cub/device/device_radix_sort.cuh>
+#include <cub/device/device_run_length_encode.cuh>
+#include <cub/util_allocator.cuh>
+
+#include <algorithm>
+#include <cstdio>
+
+#include "../../test/test_util.h"
+
+using namespace cub;
+
+//---------------------------------------------------------------------
+// Globals, constants and aliases
+//---------------------------------------------------------------------
+
+bool g_verbose = false; // Whether to display input/output to console
+CachingDeviceAllocator g_allocator(true); // Caching allocator for device memory
+
+//---------------------------------------------------------------------
+// Test generation
+//---------------------------------------------------------------------
+
+/**
+ * Simple key-value pairing for using std::sort on key-value pairs.
+ */
+template <typename Key, typename Value>
+struct Pair
+{
+  Key key;
+  Value value;
+
+  bool operator<(const Pair& b) const
+  {
+    return (key < b.key);
+  }
+};
+
+/**
+ * Pair ostream operator
+ */
+template <typename Key, typename Value>
+std::ostream& operator<<(std::ostream& os, const Pair<Key, Value>& val)
+{
+  os << '<' << val.key << ',' << val.value << '>';
+  return os;
+}
+
+/**
+ * Initialize problem
+ */
+template <typename Key, typename Value>
+void Initialize(Key* h_keys, Value* h_values, int num_items, int max_key)
+{
+  float scale = float(max_key) / float(UINT_MAX);
+  for (int i = 0; i < num_items; ++i)
+  {
+    Key sample;
+    RandomBits(sample);
+    h_keys[i]   = (max_key == -1) ? i : (Key) (scale * sample);
+    h_values[i] = i;
+  }
+
+  if (g_verbose)
+  {
+    printf("Keys:\n");
+    DisplayResults(h_keys, num_items);
+    printf("\n\n");
+
+    printf("Values:\n");
+    DisplayResults(h_values, num_items);
+    printf("\n\n");
+  }
+}
+
+/**
+ * Solve sorted non-trivial subrange problem.  Returns the number
+ * of non-trivial runs found.
+ */
+template <typename Key, typename Value>
+int Solve(Key* h_keys, Value* h_values, int num_items, int* h_offsets_reference, int* h_lengths_reference)
+{
+  // Sort
+
+  Pair<Key, Value>* h_pairs = new Pair<Key, Value>[num_items];
+  for (int i = 0; i < num_items; ++i)
+  {
+    h_pairs[i].key   = h_keys[i];
+    h_pairs[i].value = h_values[i];
+  }
+
+  std::stable_sort(h_pairs, h_pairs + num_items);
+
+  if (g_verbose)
+  {
+    printf("Sorted pairs:\n");
+    DisplayResults(h_pairs, num_items);
+    printf("\n\n");
+  }
+
+  // Find non-trivial runs
+
+  Key previous  = h_pairs[0].key;
+  int length    = 1;
+  int num_runs  = 0;
+  int run_begin = 0;
+
+  for (int i = 1; i < num_items; ++i)
+  {
+    if (previous != h_pairs[i].key)
+    {
+      if (length > 1)
+      {
+        h_offsets_reference[num_runs] = run_begin;
+        h_lengths_reference[num_runs] = length;
+        num_runs++;
+      }
+      length    = 1;
+      run_begin = i;
+    }
+    else
+    {
+      length++;
+    }
+    previous = h_pairs[i].key;
+  }
+
+  if (length > 1)
+  {
+    h_offsets_reference[num_runs] = run_begin;
+    h_lengths_reference[num_runs] = length;
+    num_runs++;
+  }
+
+  delete[] h_pairs;
+
+  return num_runs;
+}
+
+//---------------------------------------------------------------------
+// Main
+//---------------------------------------------------------------------
+
+/**
+ * Main
+ */
+int main(int argc, char** argv)
+{
+  using Key   = unsigned int;
+  using Value = int;
+
+  int timing_iterations = 0;
+  int num_items         = 40;
+  Key max_key           = 20; // Max item
+
+  // Initialize command line
+  CommandLineArgs args(argc, argv);
+  g_verbose = args.CheckCmdLineFlag("v");
+  args.GetCmdLineArgument("n", num_items);
+  args.GetCmdLineArgument("maxkey", max_key);
+  args.GetCmdLineArgument("i", timing_iterations);
+
+  // Print usage
+  if (args.CheckCmdLineFlag("help"))
+  {
+    printf("%s "
+           "[--device=<device-id>] "
+           "[--i=<timing iterations> "
+           "[--n=<input items, default 40> "
+           "[--maxkey=<max key, default 20 (use -1 to test only unique keys)>]"
+           "[--v] "
+           "\n",
+           argv[0]);
+    exit(0);
+  }
+
+  // Initialize device
+  CubDebugExit(args.DeviceInit());
+
+  // Allocate host arrays (problem and reference solution)
+
+  Key* h_keys              = new Key[num_items];
+  Value* h_values          = new Value[num_items];
+  int* h_offsets_reference = new int[num_items];
+  int* h_lengths_reference = new int[num_items];
+
+  // Initialize key-value pairs and compute reference solution (sort them, and identify non-trivial runs)
+  printf("Computing reference solution on CPU for %d items (max key %d)\n", num_items, max_key);
+  fflush(stdout);
+
+  Initialize(h_keys, h_values, num_items, static_cast<int>(max_key));
+  int num_runs = Solve(h_keys, h_values, num_items, h_offsets_reference, h_lengths_reference);
+
+  printf("%d non-trivial runs\n", num_runs);
+  fflush(stdout);
+
+  // Repeat for performance timing
+  GpuTimer gpu_timer;
+  GpuTimer gpu_rle_timer;
+  float elapsed_millis     = 0.0;
+  float elapsed_rle_millis = 0.0;
+  for (int i = 0; i <= timing_iterations; ++i)
+  {
+    // Allocate and initialize device arrays for sorting
+    DoubleBuffer<Key> d_keys;
+    DoubleBuffer<Value> d_values;
+    CubDebugExit(g_allocator.DeviceAllocate((void**) &d_keys.d_buffers[0], sizeof(Key) * num_items));
+    CubDebugExit(g_allocator.DeviceAllocate((void**) &d_keys.d_buffers[1], sizeof(Key) * num_items));
+    CubDebugExit(g_allocator.DeviceAllocate((void**) &d_values.d_buffers[0], sizeof(Value) * num_items));
+    CubDebugExit(g_allocator.DeviceAllocate((void**) &d_values.d_buffers[1], sizeof(Value) * num_items));
+
+    CubDebugExit(
+      cudaMemcpy(d_keys.d_buffers[d_keys.selector], h_keys, sizeof(float) * num_items, cudaMemcpyHostToDevice));
+    CubDebugExit(
+      cudaMemcpy(d_values.d_buffers[d_values.selector], h_values, sizeof(int) * num_items, cudaMemcpyHostToDevice));
+
+    // Start timer
+    gpu_timer.Start();
+
+    // Allocate temporary storage for sorting
+    size_t temp_storage_bytes = 0;
+    void* d_temp_storage      = nullptr;
+    CubDebugExit(DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items));
+    CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));
+
+    // Do the sort
+    CubDebugExit(DeviceRadixSort::SortPairs(d_temp_storage, temp_storage_bytes, d_keys, d_values, num_items));
+
+    // Free unused buffers and sorting temporary storage
+    if (d_keys.d_buffers[d_keys.selector ^ 1])
+    {
+      CubDebugExit(g_allocator.DeviceFree(d_keys.d_buffers[d_keys.selector ^ 1]));
+    }
+    if (d_values.d_buffers[d_values.selector ^ 1])
+    {
+      CubDebugExit(g_allocator.DeviceFree(d_values.d_buffers[d_values.selector ^ 1]));
+    }
+    if (d_temp_storage)
+    {
+      CubDebugExit(g_allocator.DeviceFree(d_temp_storage));
+    }
+
+    // Start timer
+    gpu_rle_timer.Start();
+
+    // Allocate device arrays for enumerating non-trivial runs
+    int* d_offests_out = nullptr;
+    int* d_lengths_out = nullptr;
+    int* d_num_runs    = nullptr;
+    CubDebugExit(g_allocator.DeviceAllocate((void**) &d_offests_out, sizeof(int) * num_items));
+    CubDebugExit(g_allocator.DeviceAllocate((void**) &d_lengths_out, sizeof(int) * num_items));
+    CubDebugExit(g_allocator.DeviceAllocate((void**) &d_num_runs, sizeof(int) * 1));
+
+    // Allocate temporary storage for isolating non-trivial runs
+    d_temp_storage = nullptr;
+    CubDebugExit(DeviceRunLengthEncode::NonTrivialRuns(
+      d_temp_storage,
+      temp_storage_bytes,
+      d_keys.d_buffers[d_keys.selector],
+      d_offests_out,
+      d_lengths_out,
+      d_num_runs,
+      num_items));
+    CubDebugExit(g_allocator.DeviceAllocate(&d_temp_storage, temp_storage_bytes));
+
+    // Do the isolation
+    CubDebugExit(DeviceRunLengthEncode::NonTrivialRuns(
+      d_temp_storage,
+      temp_storage_bytes,
+      d_keys.d_buffers[d_keys.selector],
+      d_offests_out,
+      d_lengths_out,
+      d_num_runs,
+      num_items));
+
+    // Free keys buffer
+    if (d_keys.d_buffers[d_keys.selector])
+    {
+      CubDebugExit(g_allocator.DeviceFree(d_keys.d_buffers[d_keys.selector]));
+    }
+
+    //
+    // Hypothetically do stuff with the original key-indices corresponding to non-trivial runs of identical keys
+    //
+
+    // Stop sort timer
+    gpu_timer.Stop();
+    gpu_rle_timer.Stop();
+
+    if (i == 0)
+    {
+      // First iteration is a warmup: // Check for correctness (and display results, if specified)
+
+      printf("\nRUN OFFSETS: \n");
+      int compare = CompareDeviceResults(h_offsets_reference, d_offests_out, num_runs, true, g_verbose);
+      printf("\t\t %s ", compare ? "FAIL" : "PASS");
+
+      printf("\nRUN LENGTHS: \n");
+      compare |= CompareDeviceResults(h_lengths_reference, d_lengths_out, num_runs, true, g_verbose);
+      printf("\t\t %s ", compare ? "FAIL" : "PASS");
+
+      printf("\nNUM RUNS: \n");
+      compare |= CompareDeviceResults(&num_runs, d_num_runs, 1, true, g_verbose);
+      printf("\t\t %s ", compare ? "FAIL" : "PASS");
+
+      AssertEquals(0, compare);
+    }
+    else
+    {
+      elapsed_millis += gpu_timer.ElapsedMillis();
+      elapsed_rle_millis += gpu_rle_timer.ElapsedMillis();
+    }
+
+    // GPU cleanup
+
+    if (d_values.d_buffers[d_values.selector])
+    {
+      CubDebugExit(g_allocator.DeviceFree(d_values.d_buffers[d_values.selector]));
+    }
+    if (d_offests_out)
+    {
+      CubDebugExit(g_allocator.DeviceFree(d_offests_out));
+    }
+    if (d_lengths_out)
+    {
+      CubDebugExit(g_allocator.DeviceFree(d_lengths_out));
+    }
+    if (d_num_runs)
+    {
+      CubDebugExit(g_allocator.DeviceFree(d_num_runs));
+    }
+    if (d_temp_storage)
+    {
+      CubDebugExit(g_allocator.DeviceFree(d_temp_storage));
+    }
+  }
+
+  // Host cleanup
+  if (h_keys)
+  {
+    delete[] h_keys;
+  }
+  if (h_values)
+  {
+    delete[] h_values;
+  }
+  if (h_offsets_reference)
+  {
+    delete[] h_offsets_reference;
+  }
+  if (h_lengths_reference)
+  {
+    delete[] h_lengths_reference;
+  }
+
+  printf("\n\n");
+
+  if (timing_iterations > 0)
+  {
+    printf("%d timing iterations, average time to sort and isolate non-trivial duplicates: %.3f ms (%.3f ms spent in "
+           "RLE isolation)\n",
+           timing_iterations,
+           elapsed_millis / static_cast<float>(timing_iterations),
+           elapsed_rle_millis / static_cast<float>(timing_iterations));
+  }
+
+  return 0;
+}

cccl_upstream/cub/examples/device/example_device_topk_keys.cu

@@ -0,0 +1,147 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+//! Simple example of cub::DeviceTopK::MinKeys().
+//! Find the top-k smallest float keys paired with a corresponding array of int values.
+//! To compile using the command line:
+//!   nvcc -arch=sm_XX example_device_topk_keys.cu -I../.. -lcudart -O3
+
+// Ensure printing of CUDA runtime errors to console
+#define CUB_STDERR
+
+#include <cub/device/device_topk.cuh>
+#include <cub/util_allocator.cuh>
+
+#include <thrust/detail/raw_pointer_cast.h>
+#include <thrust/device_vector.h>
+#include <thrust/host_vector.h>
+#include <thrust/sort.h>
+
+#include <cuda/stream>
+
+#include <algorithm>
+#include <iostream>
+
+#include "../../test/test_util.h"
+
+using namespace cub;
+
+//---------------------------------------------------------------------
+// Globals
+//---------------------------------------------------------------------
+
+// Whether to display input/output to console
+bool g_verbose = false;
+
+//---------------------------------------------------------------------
+// Helper functions
+//---------------------------------------------------------------------
+
+// Initialize the input data and the reference solution
+void initialize(float* h_keys, float* h_reference_keys, int num_items, int k)
+{
+  for (int i = 0; i < num_items; ++i)
+  {
+    RandomBits(h_keys[i]);
+  }
+
+  if (g_verbose)
+  {
+    std::cout << "Input keys:\n";
+    DisplayResults(h_keys, num_items);
+    std::cout << "\n\n";
+  }
+
+  std::partial_sort_copy(h_keys, h_keys + num_items, h_reference_keys, h_reference_keys + k);
+}
+
+//  In this example, we do no require a specific output order and do not require deterministic results (this allows for
+//  better performance in some cases). However, the output of DeviceTopK::MinKeys() is not sorted. This function sorts
+//  the output keys for comparison against the reference solution.
+thrust::host_vector<float> sort_unordered_results(thrust::host_vector<float> h_res_keys)
+{
+  thrust::sort(h_res_keys.begin(), h_res_keys.end());
+  return h_res_keys;
+}
+
+//---------------------------------------------------------------------
+// Main
+//--------------------------------------------------------------------
+int main(int argc, char** argv)
+{
+  int num_items = 10240;
+  int k         = 10;
+
+  // initialize command line
+  CommandLineArgs args(argc, argv);
+  g_verbose = args.CheckCmdLineFlag("v");
+  args.GetCmdLineArgument("n", num_items);
+  args.GetCmdLineArgument("k", k);
+
+  // Print usage
+  if (args.CheckCmdLineFlag("help"))
+  {
+    std::cout << "Usage: " << argv[0] << " [--n=<input items>] [--k=<output items>] [--device=<device-id>] [--v]\n";
+    exit(0);
+  }
+
+  // Initialize device
+  CubDebugExit(args.DeviceInit());
+
+  std::cout << "cub::DeviceTopK::MinKeys() find " << k << " smallest items from " << num_items << " items ("
+            << sizeof(float) << "-byte keys)\n";
+
+  // Allocate host arrays
+  thrust::host_vector<float> h_keys_vector(num_items);
+  thrust::host_vector<float> h_reference_keys_vector(k);
+
+  // Initialize problem and solution on host
+  initialize(thrust::raw_pointer_cast(h_keys_vector.data()),
+             thrust::raw_pointer_cast(h_reference_keys_vector.data()),
+             num_items,
+             k);
+
+  // Allocate device arrays
+  thrust::device_vector<float> d_keys_in{h_keys_vector};
+  thrust::device_vector<float> d_keys_out(k);
+
+  // Specify that we do not require a specific output order and do not require deterministic results
+  auto requirements =
+    cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted);
+
+  // Prepare CUDA stream
+  cudaStream_t stream = nullptr;
+  CubDebugExit(cudaStreamCreate(&stream));
+  cuda::stream_ref stream_ref{stream};
+
+  // Create the environment with the stream and requirements
+  auto env = cuda::std::execution::env{stream_ref, requirements};
+
+  // Query temporary storage requirements
+  size_t temp_storage_bytes = 0;
+  CubDebugExit(
+    DeviceTopK::MinKeys(nullptr, temp_storage_bytes, d_keys_in.begin(), d_keys_out.begin(), num_items, k, env));
+
+  // Allocate temporary storage
+  thrust::device_vector<std::uint8_t> temp_storage(temp_storage_bytes, thrust::no_init);
+  void* d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());
+
+  // Run the top-k algorithm
+  CubDebugExit(
+    DeviceTopK::MinKeys(d_temp_storage, temp_storage_bytes, d_keys_in.begin(), d_keys_out.begin(), num_items, k, env));
+
+  // Check for correctness (and display results, if specified)
+  auto h_res_keys_vector = sort_unordered_results(d_keys_out);
+  if (g_verbose)
+  {
+    std::cout << "Output keys:\n";
+    DisplayResults(thrust::raw_pointer_cast(h_res_keys_vector.data()), k);
+    std::cout << "\n\n";
+  }
+  const int compare = CompareResults(h_reference_keys_vector.data(), h_res_keys_vector.data(), k, g_verbose);
+  AssertEquals(0, compare);
+
+  std::cout << "\n\n";
+
+  return 0;
+}

cccl_upstream/cub/examples/device/example_device_topk_pairs.cu

@@ -0,0 +1,190 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+//! Simple example of cub::DeviceTopK::MinPairs().
+//! Find the top-k smallest float keys paired with a corresponding array of int values.
+//! To compile using the command line:
+//!   nvcc -arch=sm_XX example_device_topk_pairs.cu -I../.. -lcudart -O3
+
+// Ensure printing of CUDA runtime errors to console
+
+#define CUB_STDERR
+
+#include <cub/device/device_topk.cuh>
+#include <cub/util_allocator.cuh>
+
+#include <thrust/device_vector.h>
+#include <thrust/host_vector.h>
+#include <thrust/iterator/zip_iterator.h>
+#include <thrust/sort.h>
+
+#include <cuda/std/tuple>
+#include <cuda/stream>
+
+#include <algorithm>
+#include <iostream>
+
+#include "../../test/test_util.h"
+
+using namespace cub;
+
+//---------------------------------------------------------------------
+// Globals, constants and aliases
+//---------------------------------------------------------------------
+
+// Whether to display input/output to console
+bool g_verbose = false;
+
+//---------------------------------------------------------------------
+// Helper functions
+//---------------------------------------------------------------------
+
+// Initialize key-value sorting problem.
+void initialize(float* h_keys, int* h_values, float* h_reference_keys, int* h_reference_values, int num_items, int k)
+{
+  for (int i = 0; i < num_items; ++i)
+  {
+    RandomBits(h_keys[i]);
+    RandomBits(h_values[i]);
+  }
+
+  if (g_verbose)
+  {
+    std::cout << "Input keys:\n";
+    DisplayResults(h_keys, num_items);
+    std::cout << "\n\n";
+
+    std::cout << "Input values:\n";
+    DisplayResults(h_values, num_items);
+    std::cout << "\n\n";
+  }
+
+  auto h_pairs           = thrust::make_zip_iterator(h_keys, h_values);
+  auto h_reference_pairs = thrust::make_zip_iterator(h_reference_keys, h_reference_values);
+  std::partial_sort_copy(h_pairs, h_pairs + num_items, h_reference_pairs, h_reference_pairs + k);
+}
+
+//  In this example, we do no require a specific output order and do not require deterministic results (this allows for
+//  better performance in some cases). However, the output of DeviceTopK::MinPairs() is not sorted. This function sorts
+//  the output keys for comparison against the reference solution.
+::cuda::std::tuple<thrust::host_vector<float>, thrust::host_vector<int>>
+sort_unordered_results(thrust::host_vector<float> h_res_keys, thrust::host_vector<int> h_res_values)
+{
+  auto h_pairs = thrust::make_zip_iterator(h_res_keys.begin(), h_res_values.begin());
+  thrust::sort(h_pairs, h_pairs + static_cast<std::ptrdiff_t>(h_res_keys.size()));
+  return ::cuda::std::make_tuple(h_res_keys, h_res_values);
+}
+
+//---------------------------------------------------------------------
+// Main
+//---------------------------------------------------------------------
+int main(int argc, char** argv)
+{
+  int num_items = 10240;
+  int k         = 10;
+
+  // Initialize command line
+  CommandLineArgs args(argc, argv);
+  g_verbose = args.CheckCmdLineFlag("v");
+  args.GetCmdLineArgument("n", num_items);
+  args.GetCmdLineArgument("k", k);
+
+  // Print usage
+  if (args.CheckCmdLineFlag("help"))
+  {
+    std::cout << "Usage: " << argv[0] << " [--n=<input items>] [--k=<output items>] [--device=<device-id>] [--v]\n";
+    exit(0);
+  }
+
+  // Initialize device
+  CubDebugExit(args.DeviceInit());
+
+  std::cout << "cub::DeviceTopK::MinPairs() find " << k << " smallest items from " << num_items << " items ("
+            << sizeof(float) << "-byte keys " << sizeof(int) << "-byte values)\n";
+
+  // Allocate host arrays
+  thrust::host_vector<float> h_keys_vector(num_items);
+  thrust::host_vector<float> h_reference_keys_vector(k);
+  thrust::host_vector<float> h_res_keys_vector(k);
+  thrust::host_vector<int> h_values_vector(num_items);
+  thrust::host_vector<int> h_reference_values_vector(k);
+  thrust::host_vector<int> h_res_values_vector(k);
+
+  // Initialize problem and solution on host
+  initialize(h_keys_vector.data(),
+             h_values_vector.data(),
+             h_reference_keys_vector.data(),
+             h_reference_values_vector.data(),
+             num_items,
+             k);
+
+  // Allocate device arrays
+  thrust::device_vector<float> d_keys_in{h_keys_vector};
+  thrust::device_vector<int> d_values_in{h_values_vector};
+  thrust::device_vector<float> d_keys_out(k);
+  thrust::device_vector<int> d_values_out(k);
+
+  // Allocate temporary storage
+  size_t temp_storage_bytes = 0;
+
+  // Specify that we do not require a specific output order and do not require deterministic results
+  auto requirements =
+    cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted);
+
+  // Prepare CUDA stream
+  cudaStream_t stream = nullptr;
+  CubDebugExit(cudaStreamCreate(&stream));
+  cuda::stream_ref stream_ref{stream};
+
+  // Create the environment with the stream and requirements
+  auto env = cuda::std::execution::env{stream_ref, requirements};
+
+  // Query temporary storage requirements
+  CubDebugExit(DeviceTopK::MinPairs(
+    nullptr,
+    temp_storage_bytes,
+    d_keys_in.begin(),
+    d_keys_out.begin(),
+    d_values_in.begin(),
+    d_values_out.begin(),
+    num_items,
+    k,
+    env));
+
+  // Allocate temporary storage
+  thrust::device_vector<std::uint8_t> temp_storage(temp_storage_bytes, thrust::no_init);
+  void* d_temp_storage = thrust::raw_pointer_cast(temp_storage.data());
+
+  // Run the top-k algorithm
+  CubDebugExit(DeviceTopK::MinPairs(
+    d_temp_storage,
+    temp_storage_bytes,
+    d_keys_in.begin(),
+    d_keys_out.begin(),
+    d_values_in.begin(),
+    d_values_out.begin(),
+    num_items,
+    k,
+    env));
+
+  // Check for correctness (and display results, if specified)
+  auto [h_res_keys, h_res_values] = sort_unordered_results(d_keys_out, d_values_out);
+  if (g_verbose)
+  {
+    std::cout << "Output keys:\n";
+    DisplayResults(h_res_keys, k);
+    std::cout << "\n\n";
+
+    std::cout << "Output values:\n";
+    DisplayResults(h_res_values, k);
+    std::cout << "\n\n";
+  }
+  int compare = CompareResults(h_reference_keys_vector.data(), h_res_keys.data(), k, g_verbose);
+  AssertEquals(0, compare);
+  compare = CompareResults(h_reference_values_vector.data(), h_res_values.data(), k, g_verbose);
+  AssertEquals(0, compare);
+
+  std::cout << "\n\n";
+
+  return 0;
+}