[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/benchmarks/bench/topk/keys.cu

@@ -0,0 +1,111 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <cub/device/device_topk.cuh>
+
+#include <cuda/__execution/determinism.h>
+#include <cuda/__execution/output_ordering.h>
+#include <cuda/__execution/require.h>
+#include <cuda/__execution/tune.h>
+
+#include <nvbench_helper.cuh>
+
+// %RANGE% TUNE_ITEMS_PER_THREAD ipt 1:24:1
+// %RANGE% TUNE_THREADS_PER_BLOCK tpb 128:1024:32
+// %RANGE% TUNE_BLOCK_LOAD_ALGORITHM ld 0:2:1
+
+#if !TUNE_BASE
+template <class KeyInT>
+struct policy_selector_t
+{
+  [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const
+    -> cub::detail::topk::topk_policy
+  {
+#  if TUNE_BLOCK_LOAD_ALGORITHM == 0
+    constexpr auto load_alg = cub::BLOCK_LOAD_DIRECT;
+#  elif TUNE_BLOCK_LOAD_ALGORITHM == 1
+    constexpr auto load_alg = cub::BLOCK_LOAD_WARP_TRANSPOSE;
+#  elif TUNE_BLOCK_LOAD_ALGORITHM == 2
+    constexpr auto load_alg = cub::BLOCK_LOAD_VECTORIZE;
+#  endif
+
+    constexpr int nominal_4b_items_per_thread = TUNE_ITEMS_PER_THREAD;
+    constexpr int items_per_thread            = cuda::std::max(1, (nominal_4b_items_per_thread * 4 / sizeof(KeyInT)));
+    return cub::detail::topk::topk_policy{
+      TUNE_THREADS_PER_BLOCK,
+      items_per_thread,
+      load_alg,
+      cub::BLOCK_SCAN_WARP_SCANS,
+      cub::detail::topk::calc_bits_per_pass<KeyInT>()};
+  }
+};
+#endif // !TUNE_BASE
+
+template <typename KeyT, typename OffsetT, typename OutOffsetT>
+void topk_keys(nvbench::state& state, nvbench::type_list<KeyT, OffsetT, OutOffsetT>)
+{
+  // Retrieve axis parameters
+  const auto elements          = static_cast<size_t>(state.get_int64("Elements{io}"));
+  const auto selected_elements = static_cast<size_t>(state.get_int64("SelectedElements"));
+  const bit_entropy entropy    = str_to_entropy(state.get_string("Entropy"));
+
+  // Skip benchmarks at runtime
+  if (selected_elements >= elements)
+  {
+    state.skip("We only support the case where the variable SelectedElements is smaller than the variable "
+               "Elements{io}.");
+    return;
+  }
+
+  thrust::device_vector<KeyT> in_keys = generate(elements, entropy);
+  thrust::device_vector<KeyT> out_keys(selected_elements, thrust::no_init);
+  const KeyT* d_keys_in = thrust::raw_pointer_cast(in_keys.data());
+  KeyT* d_keys_out      = thrust::raw_pointer_cast(out_keys.data());
+
+  state.add_element_count(elements, "NumElements");
+  state.add_element_count(selected_elements, "NumSelectedElements");
+  state.add_global_memory_reads<KeyT>(elements, "InputKeys");
+  state.add_global_memory_writes<KeyT>(selected_elements, "OutputKeys");
+
+  // TODO(bgruber): call cub::DeviceTopK::MaxKeys with a the caching_allocator_t once we have an env-overload without
+  // temporary storage
+  auto env = cuda::std::execution::env{
+    cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)
+#if !TUNE_BASE
+      ,
+    cuda::execution::tune(policy_selector_t<KeyT>{})
+#endif // !TUNE_BASE
+  };
+
+  // Allocate temporary storage
+  size_t temp_size{};
+  cub::DeviceTopK::MaxKeys(
+    nullptr,
+    temp_size,
+    d_keys_in,
+    d_keys_out,
+    static_cast<OffsetT>(elements),
+    static_cast<OutOffsetT>(selected_elements),
+    env);
+  thrust::device_vector<nvbench::uint8_t> temp(temp_size, thrust::no_init);
+  auto* temp_storage = thrust::raw_pointer_cast(temp.data());
+
+  state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
+    auto env_with_stream = cuda::std::execution::env{cuda::stream_ref{launch.get_stream().get_stream()}, env};
+    cub::DeviceTopK::MaxKeys(
+      temp_storage,
+      temp_size,
+      d_keys_in,
+      d_keys_out,
+      static_cast<OffsetT>(elements),
+      static_cast<OutOffsetT>(selected_elements),
+      env_with_stream);
+  });
+}
+
+NVBENCH_BENCH_TYPES(topk_keys, NVBENCH_TYPE_AXES(fundamental_types, offset_types, offset_types))
+  .set_name("base")
+  .set_type_axes_names({"KeyT{ct}", "OffsetT{ct}", "OutOffsetT{ct}"})
+  .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4))
+  .add_int64_power_of_two_axis("SelectedElements", nvbench::range(3, 23, 4))
+  .add_string_axis("Entropy", {"1.000", "0.544", "0.201", "0.000"});

cccl_upstream/cub/benchmarks/bench/topk/pairs.cu

@@ -0,0 +1,120 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <cub/device/device_topk.cuh>
+
+#include <cuda/__execution/determinism.h>
+#include <cuda/__execution/output_ordering.h>
+#include <cuda/__execution/require.h>
+#include <cuda/__execution/tune.h>
+
+#include <nvbench_helper.cuh>
+
+// %RANGE% TUNE_ITEMS_PER_THREAD ipt 1:24:1
+// %RANGE% TUNE_THREADS_PER_BLOCK tpb 128:1024:32
+// %RANGE% TUNE_BLOCK_LOAD_ALGORITHM ld 0:2:1
+
+#if !TUNE_BASE
+template <class KeyInT>
+struct policy_selector_t
+{
+  [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const
+    -> cub::detail::topk::topk_policy
+  {
+#  if TUNE_BLOCK_LOAD_ALGORITHM == 0
+    constexpr auto load_alg = cub::BLOCK_LOAD_DIRECT;
+#  elif TUNE_BLOCK_LOAD_ALGORITHM == 1
+    constexpr auto load_alg = cub::BLOCK_LOAD_WARP_TRANSPOSE;
+#  elif TUNE_BLOCK_LOAD_ALGORITHM == 2
+    constexpr auto load_alg = cub::BLOCK_LOAD_VECTORIZE;
+#  endif
+
+    constexpr int nominal_4b_items_per_thread = TUNE_ITEMS_PER_THREAD;
+    constexpr int items_per_thread            = cuda::std::max(1, (nominal_4b_items_per_thread * 4 / sizeof(KeyInT)));
+    return cub::detail::topk::topk_policy{
+      TUNE_THREADS_PER_BLOCK,
+      items_per_thread,
+      load_alg,
+      cub::BLOCK_SCAN_WARP_SCANS,
+      cub::detail::topk::calc_bits_per_pass<KeyInT>()};
+  }
+};
+#endif // !TUNE_BASE
+
+template <typename KeyT, typename ValueT, typename OffsetT, typename OutOffsetT>
+void topk_pairs(nvbench::state& state, nvbench::type_list<KeyT, ValueT, OffsetT, OutOffsetT>)
+{
+  // Retrieve axis parameters
+  const auto elements          = static_cast<size_t>(state.get_int64("Elements{io}"));
+  const auto selected_elements = static_cast<size_t>(state.get_int64("SelectedElements"));
+  const bit_entropy entropy    = str_to_entropy(state.get_string("Entropy"));
+
+  // Skip benchmarks at runtime
+  if (selected_elements >= elements)
+  {
+    state.skip("We only support the case where the variable SelectedElements is smaller than the variable "
+               "Elements{io}.");
+    return;
+  }
+
+  thrust::device_vector<KeyT> in_keys     = generate(elements, entropy);
+  thrust::device_vector<ValueT> in_values = generate(elements);
+  thrust::device_vector<KeyT> out_keys(selected_elements, thrust::no_init);
+  thrust::device_vector<ValueT> out_values(selected_elements, thrust::no_init);
+
+  const KeyT* d_keys_in     = thrust::raw_pointer_cast(in_keys.data());
+  KeyT* d_keys_out          = thrust::raw_pointer_cast(out_keys.data());
+  const ValueT* d_values_in = thrust::raw_pointer_cast(in_values.data());
+  ValueT* d_values_out      = thrust::raw_pointer_cast(out_values.data());
+
+  state.add_element_count(elements, "NumElements");
+  state.add_element_count(selected_elements, "NumSelectedElements");
+  state.add_global_memory_reads<KeyT>(elements, "InputKeys");
+  state.add_global_memory_reads<ValueT>(elements, "InputValues");
+  state.add_global_memory_writes<KeyT>(selected_elements, "OutputKeys");
+  state.add_global_memory_writes<ValueT>(selected_elements, "OutputVales");
+
+  auto env = cuda::std::execution::env{
+    cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted)
+#if !TUNE_BASE
+      ,
+    cuda::execution::tune(policy_selector_t<KeyT>{})
+#endif // !TUNE_BASE
+  };
+
+  // Allocate temporary storage
+  size_t temp_size{};
+  cub::DeviceTopK::MaxPairs(
+    nullptr,
+    temp_size,
+    d_keys_in,
+    d_keys_out,
+    d_values_in,
+    d_values_out,
+    static_cast<OffsetT>(elements),
+    static_cast<OutOffsetT>(selected_elements),
+    env);
+  thrust::device_vector<nvbench::uint8_t> temp(temp_size, thrust::no_init);
+  auto* temp_storage = thrust::raw_pointer_cast(temp.data());
+
+  state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
+    auto env_with_stream = cuda::std::execution::env{cuda::stream_ref{launch.get_stream().get_stream()}, env};
+    cub::DeviceTopK::MaxPairs(
+      temp_storage,
+      temp_size,
+      d_keys_in,
+      d_keys_out,
+      d_values_in,
+      d_values_out,
+      static_cast<OffsetT>(elements),
+      static_cast<OutOffsetT>(selected_elements),
+      env_with_stream);
+  });
+}
+
+NVBENCH_BENCH_TYPES(topk_pairs, NVBENCH_TYPE_AXES(integral_types, integral_types, offset_types, offset_types))
+  .set_name("base")
+  .set_type_axes_names({"KeyT{ct}", "ValueT{ct}", "OffsetT{ct}", "OutOffsetT{ct}"})
+  .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4))
+  .add_int64_power_of_two_axis("SelectedElements", nvbench::range(3, 23, 4))
+  .add_string_axis("Entropy", {"1.000", "0.544", "0.201", "0.000"});