[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/segmented_reduce/argmin.cu

@@ -0,0 +1,16 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <nvbench_helper.cuh>
+
+// %RANGE% TUNE_ITEMS_PER_VEC_LOAD_POW2 ipv 1:2:1
+// %RANGE% TUNE_S_THREADS_PER_WARP stpw 1:32:1
+// %RANGE% TUNE_M_THREADS_PER_WARP mtpw 1:32:1
+// %RANGE% TUNE_L_NOMINAL_4B_THREADS_PER_BLOCK ltpb 128:1024:32
+// %RANGE% TUNE_S_NOMINAL_4B_ITEMS_PER_THREAD sipt 1:32:1
+// %RANGE% TUNE_M_NOMINAL_4B_ITEMS_PER_THREAD mipt 1:32:1
+// %RANGE% TUNE_L_NOMINAL_4B_ITEMS_PER_THREAD lipt 7:24:1
+
+using value_types = integral_types;
+using op_t        = cub::detail::arg_min;
+#include "base.cuh"

cccl_upstream/cub/benchmarks/bench/segmented_reduce/base.cuh

@@ -0,0 +1,122 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+#pragma once
+
+#include <cub/device/device_segmented_reduce.cuh>
+
+#include <cuda/std/type_traits>
+
+#ifndef TUNE_BASE
+#  define TUNE_ITEMS_PER_VEC_LOAD (1 << TUNE_ITEMS_PER_VEC_LOAD_POW2)
+#endif
+
+#if !TUNE_BASE
+template <typename AccumT>
+struct policy_selector
+{
+  [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(cuda::compute_capability) const
+    -> ::cub::SegmentedReducePolicy
+  {
+    constexpr int accum_size = int{sizeof(AccumT)};
+
+    const auto [l_items, l_threads] =
+      cub::detail::scale_mem_bound(TUNE_L_NOMINAL_4B_THREADS_PER_BLOCK, TUNE_L_NOMINAL_4B_ITEMS_PER_THREAD, accum_size);
+    const auto s_items =
+      cub::detail::scale_mem_bound(TUNE_L_NOMINAL_4B_THREADS_PER_BLOCK, TUNE_S_NOMINAL_4B_ITEMS_PER_THREAD, accum_size)
+        .items_per_thread;
+    const auto m_items =
+      cub::detail::scale_mem_bound(TUNE_L_NOMINAL_4B_THREADS_PER_BLOCK, TUNE_M_NOMINAL_4B_ITEMS_PER_THREAD, accum_size)
+        .items_per_thread;
+
+    const auto rp = cub::ReducePassPolicy{
+      l_threads, l_items, TUNE_ITEMS_PER_VEC_LOAD, cub::BLOCK_REDUCE_WARP_REDUCTIONS, cub::LOAD_LDG};
+    return {rp,
+            cub::SegmentedReduceWarpReducePolicy{
+              rp.threads_per_block, TUNE_M_THREADS_PER_WARP, m_items, rp.vec_size, rp.load_modifier},
+            cub::SegmentedReduceWarpReducePolicy{
+              rp.threads_per_block, TUNE_S_THREADS_PER_WARP, s_items, rp.vec_size, rp.load_modifier}};
+  }
+};
+#endif // !TUNE_BASE
+
+template <typename T>
+void fixed_size_segmented_reduce(nvbench::state& state, nvbench::type_list<T>)
+{
+  static constexpr bool is_argmin = std::is_same_v<op_t, cub::detail::arg_min>;
+
+  using output_t     = cuda::std::conditional_t<is_argmin, cuda::std::pair<int, T>, T>;
+  using accum_t      = output_t;
+  using init_value_t = cuda::std::conditional_t<is_argmin, cub::detail::reduce::empty_problem_init_t<accum_t>, T>;
+
+  // Retrieve axis parameters
+  const size_t num_elements = static_cast<size_t>(state.get_int64("Elements{io}"));
+  const size_t segment_size = static_cast<size_t>(state.get_int64("SegmentSize"));
+  const size_t num_segments = std::max<std::size_t>(1, (num_elements / segment_size));
+  const size_t elements     = num_segments * segment_size;
+
+  thrust::device_vector<T> in = generate(elements);
+  thrust::device_vector<output_t> out(num_segments);
+
+  const T* d_in   = thrust::raw_pointer_cast(in.data());
+  output_t* d_out = thrust::raw_pointer_cast(out.data());
+
+  // Enable throughput calculations and add "Size" column to results.
+  state.add_element_count(elements);
+  state.add_global_memory_reads<T>(elements, "Size");
+  state.add_global_memory_writes<output_t>(num_segments);
+
+  caching_allocator_t alloc;
+  state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
+    auto env = cub_bench_env(
+      alloc,
+      launch
+#if !TUNE_BASE
+      ,
+      cuda::execution::tune(policy_selector<accum_t>{})
+#endif
+    );
+    if constexpr (is_argmin)
+    {
+      _CCCL_TRY_CUDA_API(
+        cub::DeviceSegmentedReduce::ArgMin,
+        "Segmented ArgMin failed",
+        d_in,
+        d_out,
+        static_cast<::cuda::std::int64_t>(num_segments),
+        static_cast<int>(segment_size),
+        env);
+    }
+    else
+    {
+      _CCCL_TRY_CUDA_API(
+        cub::DeviceSegmentedReduce::Reduce,
+        "Segmented reduce failed",
+        d_in,
+        d_out,
+        static_cast<::cuda::std::int64_t>(num_segments),
+        static_cast<int>(segment_size),
+        op_t{},
+        init_value_t{},
+        env);
+    }
+  });
+}
+
+NVBENCH_BENCH_TYPES(fixed_size_segmented_reduce, NVBENCH_TYPE_AXES(value_types))
+  .set_name("small")
+  .set_type_axes_names({"T{ct}"})
+  .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4))
+  .add_int64_power_of_two_axis("SegmentSize", nvbench::range(0, 4, 1));
+
+NVBENCH_BENCH_TYPES(fixed_size_segmented_reduce, NVBENCH_TYPE_AXES(value_types))
+  .set_name("medium")
+  .set_type_axes_names({"T{ct}"})
+  .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4))
+  .add_int64_power_of_two_axis("SegmentSize", nvbench::range(5, 8, 1));
+
+NVBENCH_BENCH_TYPES(fixed_size_segmented_reduce, NVBENCH_TYPE_AXES(value_types))
+  .set_name("large")
+  .set_type_axes_names({"T{ct}"})
+  .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4))
+  .add_int64_power_of_two_axis("SegmentSize", nvbench::range(9, 16, 1));

cccl_upstream/cub/benchmarks/bench/segmented_reduce/custom.cu

@@ -0,0 +1,8 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+#include <nvbench_helper.cuh>
+
+using value_types = all_types;
+using op_t        = max_t;
+#include "base.cuh"

cccl_upstream/cub/benchmarks/bench/segmented_reduce/sum.cu

@@ -0,0 +1,16 @@
+// SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+#include <nvbench_helper.cuh>
+
+// %RANGE% TUNE_ITEMS_PER_VEC_LOAD_POW2 ipv 1:2:1
+// %RANGE% TUNE_S_THREADS_PER_WARP stpw 1:32:1
+// %RANGE% TUNE_M_THREADS_PER_WARP mtpw 1:32:1
+// %RANGE% TUNE_L_NOMINAL_4B_THREADS_PER_BLOCK ltpb 128:1024:32
+// %RANGE% TUNE_S_NOMINAL_4B_ITEMS_PER_THREAD sipt 1:32:1
+// %RANGE% TUNE_M_NOMINAL_4B_ITEMS_PER_THREAD mipt 1:32:1
+// %RANGE% TUNE_L_NOMINAL_4B_ITEMS_PER_THREAD lipt 7:24:1
+
+using value_types = all_types;
+using op_t        = ::cuda::std::plus<>;
+#include "base.cuh"

cccl_upstream/cub/benchmarks/bench/segmented_reduce/variable_argmax.cu

@@ -0,0 +1,8 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <nvbench_helper.cuh>
+
+using op_t = cub::detail::arg_max;
+
+#include "variable_base.cuh"

cccl_upstream/cub/benchmarks/bench/segmented_reduce/variable_base.cuh

@@ -0,0 +1,161 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#pragma once
+
+#include <cub/device/dispatch/dispatch_segmented_reduce.cuh>
+
+#include <cuda/std/iterator>
+#include <cuda/std/type_traits>
+
+#include <nvbench_helper.cuh>
+
+#if TUNE_T
+using value_types = nvbench::type_list<TUNE_T>;
+#else
+using value_types = nvbench::type_list<int32_t, int64_t, float, double>;
+#endif
+
+#ifdef TUNE_OffsetT
+using some_offset_types = nvbench::type_list<TUNE_OffsetT>;
+#else
+using some_offset_types = nvbench::type_list<int32_t>;
+#endif
+
+template <typename T, typename OffsetT>
+void variable_segmented_reduce(nvbench::state& state, nvbench::type_list<T, OffsetT>)
+{
+  static constexpr bool is_argmin = std::is_same_v<op_t, cub::detail::arg_min>;
+  static constexpr bool is_argmax = std::is_same_v<op_t, cub::detail::arg_max>;
+
+  using raw_input_it_t = const T*;
+  using output_t       = cuda::std::conditional_t<(is_argmin || is_argmax), cuda::std::pair<int, T>, T>;
+  using output_it_t    = output_t*;
+  using accum_t        = output_t;
+  using init_value_t =
+    cuda::std::conditional_t<(is_argmin || is_argmax), cub::detail::reduce::empty_problem_init_t<accum_t>, T>;
+  using offset_t          = OffsetT;
+  using begin_offset_it_t = const offset_t*;
+  using end_offset_it_t   = const offset_t*;
+
+  // Retrieve axis parameters
+  const auto elements                = static_cast<std::size_t>(state.get_int64("Elements{io}"));
+  const auto max_segment_size        = static_cast<std::size_t>(state.get_int64("MaxSegmentSize"));
+  const auto guaranteed_max_seg_size = static_cast<std::size_t>(state.get_int64("GuaranteedMaxSegSize"));
+
+  // skip if max_segment_size > guaranteed_max_seg_size
+  if (guaranteed_max_seg_size != 0 && max_segment_size > guaranteed_max_seg_size)
+  {
+    state.skip("max_segment_size > guaranteed_max_seg_size");
+    return;
+  }
+
+  const auto min_segment_size     = 1;
+  const auto max_segment_size_log = static_cast<offset_t>(std::log2(max_segment_size));
+
+  // Generate segment offsets
+  thrust::device_vector<offset_t> segment_offsets =
+    generate.uniform.segment_offsets(elements, min_segment_size, max_segment_size);
+  const auto num_segments = segment_offsets.size() - 1;
+
+  // Generate input data
+  thrust::device_vector<T> in = generate(elements);
+
+  thrust::device_vector<output_t> out(num_segments, thrust::default_init);
+
+  raw_input_it_t d_raw_in           = thrust::raw_pointer_cast(in.data());
+  output_it_t d_out                 = thrust::raw_pointer_cast(out.data());
+  begin_offset_it_t d_begin_offsets = thrust::raw_pointer_cast(segment_offsets.data());
+  end_offset_it_t d_end_offsets     = d_begin_offsets + 1;
+
+  // Create wrapped iterator for argmin/argmax operations
+  [[maybe_unused]] auto d_indexed_in = cuda::make_transform_iterator(
+    cuda::counting_iterator<::cuda::std::int64_t>(0),
+    cub::detail::segmented_reduce::generate_idx_value<raw_input_it_t, T>(d_raw_in, 1));
+  using arg_index_input_iterator_t = decltype(d_indexed_in);
+
+  auto d_in = [&] {
+    if constexpr (is_argmin || is_argmax)
+    {
+      return d_indexed_in;
+    }
+    else
+    {
+      return d_raw_in;
+    }
+  }();
+
+  // Enable throughput calculations
+  state.add_element_count(elements);
+  state.add_global_memory_reads<T>(elements, "Size");
+  state.add_global_memory_writes<output_t>(num_segments);
+  state.add_global_memory_reads<offset_t>(num_segments + 1);
+
+  // Allocate temporary storage
+  std::size_t temp_size{};
+  using override_offset_t = cuda::std::conditional_t<(is_argmin || is_argmax), int, cub::detail::use_default>;
+
+  // TODO(bgruber): rewrite this to use the public CUB API directly. But in order to do this, we need to expose the
+  // guaranteed_max_seg_size at the public API
+  cub::detail::segmented_reduce::dispatch<accum_t, override_offset_t>(
+    nullptr,
+    temp_size,
+    d_in,
+    d_out,
+    static_cast<::cuda::std::int64_t>(num_segments),
+    d_begin_offsets,
+    d_end_offsets,
+    op_t{},
+    init_value_t{},
+    guaranteed_max_seg_size,
+    nullptr /* stream */);
+
+  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) {
+    cub::detail::segmented_reduce::dispatch<accum_t, override_offset_t>(
+      temp_storage,
+      temp_size,
+      d_in,
+      d_out,
+      static_cast<::cuda::std::int64_t>(num_segments),
+      d_begin_offsets,
+      d_end_offsets,
+      op_t{},
+      init_value_t{},
+      guaranteed_max_seg_size,
+      launch.get_stream());
+  });
+}
+
+NVBENCH_BENCH_TYPES(variable_segmented_reduce, NVBENCH_TYPE_AXES(value_types, some_offset_types))
+  .set_name("variable_default")
+  .set_type_axes_names({"T{ct}", "OffsetT{ct}"})
+  .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4))
+  .add_int64_power_of_two_axis("MaxSegmentSize", nvbench::range(1, 16, 1))
+  .add_int64_axis("GuaranteedMaxSegSize", {0});
+
+// Small segments: 1-16 items per segment
+NVBENCH_BENCH_TYPES(variable_segmented_reduce, NVBENCH_TYPE_AXES(value_types, some_offset_types))
+  .set_name("variable_small_dynamic")
+  .set_type_axes_names({"T{ct}", "OffsetT{ct}"})
+  .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4))
+  .add_int64_power_of_two_axis("MaxSegmentSize", nvbench::range(1, 4, 1))
+  .add_int64_power_of_two_axis("GuaranteedMaxSegSize", nvbench::range(1, 4, 1));
+
+// Medium segments: 32-256 items per segment
+NVBENCH_BENCH_TYPES(variable_segmented_reduce, NVBENCH_TYPE_AXES(value_types, some_offset_types))
+  .set_name("variable_medium_dynamic")
+  .set_type_axes_names({"T{ct}", "OffsetT{ct}"})
+  .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4))
+  .add_int64_power_of_two_axis("MaxSegmentSize", nvbench::range(5, 8, 1))
+  .add_int64_power_of_two_axis("GuaranteedMaxSegSize", nvbench::range(5, 8, 1));
+
+// Large segments: 512+ items per segment
+NVBENCH_BENCH_TYPES(variable_segmented_reduce, NVBENCH_TYPE_AXES(value_types, some_offset_types))
+  .set_name("variable_large_dynamic")
+  .set_type_axes_names({"T{ct}", "OffsetT{ct}"})
+  .add_int64_power_of_two_axis("Elements{io}", nvbench::range(16, 28, 4))
+  .add_int64_power_of_two_axis("MaxSegmentSize", nvbench::range(9, 16, 1))
+  .add_int64_power_of_two_axis("GuaranteedMaxSegSize", nvbench::range(9, 16, 1));

cccl_upstream/cub/benchmarks/bench/segmented_reduce/variable_sum.cu

@@ -0,0 +1,8 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <nvbench_helper.cuh>
+
+using op_t = ::cuda::std::plus<>;
+
+#include "variable_base.cuh"