[CUB] Use the vectorized load machinery in AgentFind

cub/benchmarks/bench/load/block_load.cu

@@ -0,0 +1,136 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <cub/block/block_load.cuh>
+
+#include <thrust/device_vector.h>
+
+#include <cuda/iterator>
+#include <cuda/std/cstdint>
+
+#include <nvbench_helper.cuh>
+
+struct increment
+{
+  template <class T>
+  _CCCL_DEVICE T operator()(T value) const
+  {
+    return static_cast<T>(value + T{1});
+  }
+};
+
+struct twice
+{
+  template <class T>
+  _CCCL_DEVICE T operator()(T value) const
+  {
+    return static_cast<T>(value * T{2});
+  }
+};
+
+// Load BaseT storage through a nested transform and keep a per-thread sum. The sum is written once per block so the
+// transformed loads cannot be discarded.
+template <class BaseT, int BlockSize, int ItemsPerThread, cub::BlockLoadAlgorithm Algorithm>
+__global__ void
+block_load_transform(const BaseT* input, ::cuda::std::int64_t num_tiles, ::cuda::std::uint64_t* block_sums)
+{
+  constexpr int tile_size = BlockSize * ItemsPerThread;
+  using block_load_t      = cub::BlockLoad<BaseT, BlockSize, ItemsPerThread, Algorithm>;
+  using storage_t         = typename block_load_t::TempStorage;
+
+  __shared__ storage_t storage;
+  ::cuda::std::uint64_t sum = 0;
+  for (::cuda::std::int64_t tile = blockIdx.x; tile < num_tiles; tile += gridDim.x)
+  {
+    const BaseT* tile_input = input + tile * tile_size;
+    auto inner              = ::cuda::transform_iterator<increment, const BaseT*>{tile_input, increment{}};
+    auto in                 = ::cuda::transform_iterator<twice, decltype(inner)>{inner, twice{}};
+
+    BaseT data[ItemsPerThread];
+    block_load_t{storage}.Load(in, data);
+    __syncthreads();
+
+    _CCCL_PRAGMA_UNROLL_FULL()
+    for (int item = 0; item < ItemsPerThread; ++item)
+    {
+      sum += static_cast<::cuda::std::uint64_t>(data[item]);
+    }
+  }
+
+  if (threadIdx.x == 0)
+  {
+    block_sums[blockIdx.x] = sum;
+  }
+}
+
+template <class BaseT, int BlockSize, int ItemsPerThread, cub::BlockLoadAlgorithm Algorithm>
+void run_block_load(nvbench::state& state)
+{
+  constexpr int tile_size = BlockSize * ItemsPerThread;
+  const auto requested    = state.get_int64("Elements{io}");
+  const auto num_tiles    = requested / tile_size;
+  if (num_tiles == 0)
+  {
+    state.skip("Skipping: element count is smaller than one tile.");
+    return;
+  }
+  const auto num_items = num_tiles * tile_size;
+
+  thrust::device_vector<BaseT> input(static_cast<std::size_t>(num_items));
+  const auto kernel     = block_load_transform<BaseT, BlockSize, ItemsPerThread, Algorithm>;
+  const int num_sms     = state.get_device().value().get_number_of_sms(); // NOLINT(bugprone-unchecked-optional-access)
+  int max_blocks_per_sm = 0;
+  NVBENCH_CUDA_CALL_NOEXCEPT(cudaOccupancyMaxActiveBlocksPerMultiprocessor(&max_blocks_per_sm, kernel, BlockSize, 0));
+  if (max_blocks_per_sm == 0)
+  {
+    state.skip("Skipping: no resident blocks for this load on this device.");
+    return;
+  }
+  const int grid_size = max_blocks_per_sm * num_sms;
+  thrust::device_vector<::cuda::std::uint64_t> block_sums(static_cast<std::size_t>(grid_size));
+
+  state.add_element_count(num_items);
+  state.add_global_memory_reads<BaseT>(num_items);
+  state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
+    kernel<<<grid_size, BlockSize, 0, launch.get_stream()>>>(
+      thrust::raw_pointer_cast(input.data()), num_tiles, thrust::raw_pointer_cast(block_sums.data()));
+  });
+}
+
+template <class BaseT>
+void block_load(nvbench::state& state, nvbench::type_list<BaseT>)
+{
+  const bool vectorize        = state.get_string("Algorithm") == "vectorize";
+  const auto items_per_thread = state.get_int64("ItemsPerThread");
+  const auto block_size       = state.get_int64("BlockSize");
+
+  if (block_size == 128 && items_per_thread == 4)
+  {
+    return vectorize ? run_block_load<BaseT, 128, 4, cub::BLOCK_LOAD_VECTORIZE>(state)
+                     : run_block_load<BaseT, 128, 4, cub::BLOCK_LOAD_DIRECT>(state);
+  }
+  if (block_size == 128 && items_per_thread == 8)
+  {
+    return vectorize ? run_block_load<BaseT, 128, 8, cub::BLOCK_LOAD_VECTORIZE>(state)
+                     : run_block_load<BaseT, 128, 8, cub::BLOCK_LOAD_DIRECT>(state);
+  }
+  if (block_size == 256 && items_per_thread == 4)
+  {
+    return vectorize ? run_block_load<BaseT, 256, 4, cub::BLOCK_LOAD_VECTORIZE>(state)
+                     : run_block_load<BaseT, 256, 4, cub::BLOCK_LOAD_DIRECT>(state);
+  }
+  if (block_size == 256 && items_per_thread == 8)
+  {
+    return vectorize ? run_block_load<BaseT, 256, 8, cub::BLOCK_LOAD_VECTORIZE>(state)
+                     : run_block_load<BaseT, 256, 8, cub::BLOCK_LOAD_DIRECT>(state);
+  }
+  state.skip("Skipping: unsupported block size or items per thread.");
+}
+
+NVBENCH_BENCH_TYPES(block_load, NVBENCH_TYPE_AXES(integral_types))
+  .set_name("transform")
+  .set_type_axes_names({"BaseT{ct}"})
+  .add_string_axis("Algorithm", {"vectorize", "direct"})
+  .add_int64_axis("BlockSize", {128, 256})
+  .add_int64_axis("ItemsPerThread", {4, 8})
+  .add_int64_power_of_two_axis("Elements{io}", nvbench::range(20, 28, 4));

cub/benchmarks/bench/reduce/contiguous_adaptor.cu

@@ -0,0 +1,65 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+#include <cub/device/device_reduce.cuh>
+
+#include <cuda/iterator>
+
+#include <nvbench_helper.cuh>
+
+// Same-type transform over a device pointer. The stored values are contiguous; the iterator itself is not.
+template <typename T>
+struct plus_one_t
+{
+  __host__ __device__ T operator()(T value) const
+  {
+    return static_cast<T>(value + T{1});
+  }
+};
+
+template <typename T>
+void sum(nvbench::state& state, nvbench::type_list<T>)
+try
+{
+  const auto elements     = state.get_int64("Elements{io}");
+  const auto source       = state.get_string("Source");
+  const bool as_transform = source == "transform";
+
+  thrust::device_vector<T> in = generate(static_cast<std::size_t>(elements));
+  thrust::device_vector<T> out(1);
+
+  auto* d_in  = thrust::raw_pointer_cast(in.data());
+  auto* d_out = thrust::raw_pointer_cast(out.data());
+
+  state.add_element_count(elements);
+  state.add_global_memory_reads<T>(elements, "Size");
+  state.add_global_memory_writes<T>(1);
+
+  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 (as_transform)
+    {
+      auto it = cuda::transform_iterator(d_in, plus_one_t<T>{});
+      _CCCL_TRY_RUNTIME_API(
+        cub::DeviceReduce::Sum, "Reduce failed", it, d_out, static_cast<std::int32_t>(elements), env);
+    }
+    else
+    {
+      _CCCL_TRY_RUNTIME_API(
+        cub::DeviceReduce::Sum, "Reduce failed", d_in, d_out, static_cast<std::int32_t>(elements), env);
+    }
+  });
+}
+catch (const std::bad_alloc&)
+{
+  state.skip("Skipping: out of memory.");
+}
+
+using value_types = nvbench::type_list<std::int8_t, std::int16_t, std::int32_t, std::int64_t, int128_t>;
+
+NVBENCH_BENCH_TYPES(sum, NVBENCH_TYPE_AXES(value_types))
+  .set_name("contiguous_adaptor")
+  .set_type_axes_names({"T{ct}"})
+  .add_string_axis("Source", {"pointer", "transform"})
+  .add_int64_power_of_two_axis("Elements{io}", nvbench::range(24, 28, 4));

cub/benchmarks/bench/topk/transform_iterator.cu

@@ -0,0 +1,96 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+
+#include <cub/device/device_topk.cuh>
+#include <cub/util_type.cuh>
+
+#include <cuda/__execution/determinism.h>
+#include <cuda/__execution/output_ordering.h>
+#include <cuda/__execution/require.h>
+#include <cuda/__execution/tune.h>
+#include <cuda/iterator>
+#include <cuda/std/cstdint>
+
+#include <nvbench_helper.cuh>
+
+// DeviceTopK is the device algorithm whose shipped tuning selects BLOCK_LOAD_VECTORIZE. The input is a
+// transform_iterator over const char*, so the vectorized load copies the underlying bytes and the functor runs when
+// the loaded items are read back.
+struct widen
+{
+  [[nodiscard]] _CCCL_HOST_DEVICE uint32_t operator()(char value) const
+  {
+    return static_cast<uint32_t>(static_cast<unsigned char>(value));
+  }
+};
+
+// Same tile shape as the shipped uint32 policy, with the load forced to BLOCK_LOAD_DIRECT.
+struct direct_load_policy
+{
+  [[nodiscard]] _CCCL_HOST_DEVICE constexpr auto operator()(::cuda::compute_capability cc) const
+  {
+    auto policy = cub::detail::topk::policy_selector_from_types<uint32_t, cub::NullType, int64_t, int>{}(cc);
+    return cub::detail::topk::topk_policy{
+      policy.threads_per_block,
+      policy.items_per_thread,
+      cub::BLOCK_LOAD_DIRECT,
+      policy.scan_algorithm,
+      policy.bits_per_pass};
+  }
+};
+
+template <class Env>
+void max_keys(nvbench::state& state, int64_t elements, int selected_elements, Env env)
+{
+  thrust::device_vector<char> input(static_cast<size_t>(elements), thrust::no_init);
+  thrust::device_vector<uint32_t> output(static_cast<size_t>(selected_elements), thrust::no_init);
+
+  const auto in = cuda::transform_iterator<widen, const char*>{thrust::raw_pointer_cast(input.data()), widen{}};
+  auto* out     = thrust::raw_pointer_cast(output.data());
+
+  state.add_element_count(elements, "NumElements");
+  state.add_element_count(selected_elements, "NumSelectedElements");
+  // Traffic is the underlying char storage, not the uint32 key the functor returns.
+  state.add_global_memory_reads<char>(elements, "InputKeys");
+  state.add_global_memory_writes<uint32_t>(selected_elements, "OutputKeys");
+
+  size_t temp_size{};
+  cub::DeviceTopK::MaxKeys(nullptr, temp_size, in, out, elements, 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, in, out, elements, selected_elements, env_with_stream);
+  });
+}
+
+void topk_transform(nvbench::state& state)
+{
+  const auto elements    = state.get_int64("Elements{io}");
+  constexpr int selected = 1024;
+  const bool vectorize   = state.get_string("Load") == "vectorize";
+
+  if (selected >= elements)
+  {
+    state.skip("Selected count must be smaller than the input.");
+    return;
+  }
+
+  auto requirements =
+    cuda::execution::require(cuda::execution::determinism::not_guaranteed, cuda::execution::output_ordering::unsorted);
+  if (vectorize)
+  {
+    max_keys(state, elements, selected, cuda::std::execution::env{requirements});
+  }
+  else
+  {
+    max_keys(
+      state, elements, selected, cuda::std::execution::env{requirements, cuda::execution::tune(direct_load_policy{})});
+  }
+}
+
+NVBENCH_BENCH(topk_transform)
+  .set_name("transform")
+  .add_string_axis("Load", {"vectorize", "direct"})
+  .add_int64_power_of_two_axis("Elements{io}", nvbench::range(20, 28, 4));

cub/benchmarks/bench/transform/contiguous_adaptor.cu

@@ -0,0 +1,77 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+#include <cub/device/device_transform.cuh>
+
+#include <cuda/iterator>
+#include <cuda/std/functional>
+
+#include <nvbench_helper.cuh>
+
+// Same-type transform over a device pointer. The stored values are contiguous; the iterator itself is not.
+template <typename T>
+struct plus_one_t
+{
+  __host__ __device__ T operator()(T value) const
+  {
+    return static_cast<T>(value + T{1});
+  }
+};
+
+template <typename T>
+void copy(nvbench::state& state, nvbench::type_list<T>)
+try
+{
+  const auto elements     = state.get_int64("Elements{io}");
+  const auto source       = state.get_string("Source");
+  const bool as_transform = source == "transform";
+
+  thrust::device_vector<T> in  = generate(static_cast<std::size_t>(elements));
+  thrust::device_vector<T> out = generate(static_cast<std::size_t>(elements));
+
+  auto* d_in  = thrust::raw_pointer_cast(in.data());
+  auto* d_out = thrust::raw_pointer_cast(out.data());
+
+  state.add_element_count(elements);
+  state.add_global_memory_reads<T>(elements, "Size");
+  state.add_global_memory_writes<T>(elements);
+
+  state.exec(nvbench::exec_tag::gpu | nvbench::exec_tag::no_batch, [&](nvbench::launch& launch) {
+    auto env = cuda::std::execution::env{cuda::stream_ref{launch.get_stream().get_stream()}};
+    if (as_transform)
+    {
+      auto it = cuda::transform_iterator(d_in, plus_one_t<T>{});
+      _CCCL_TRY_RUNTIME_API(
+        cub::DeviceTransform::Transform,
+        "Transform failed",
+        cuda::std::make_tuple(it),
+        d_out,
+        static_cast<std::int32_t>(elements),
+        cuda::std::identity{},
+        env);
+    }
+    else
+    {
+      _CCCL_TRY_RUNTIME_API(
+        cub::DeviceTransform::Transform,
+        "Transform failed",
+        cuda::std::make_tuple(d_in),
+        d_out,
+        static_cast<std::int32_t>(elements),
+        plus_one_t<T>{},
+        env);
+    }
+  });
+}
+catch (const std::bad_alloc&)
+{
+  state.skip("Skipping: out of memory.");
+}
+
+using value_types = nvbench::type_list<std::int8_t, std::int16_t, std::int32_t, std::int64_t, int128_t>;
+
+NVBENCH_BENCH_TYPES(copy, NVBENCH_TYPE_AXES(value_types))
+  .set_name("contiguous_adaptor")
+  .set_type_axes_names({"T{ct}"})
+  .add_string_axis("Source", {"pointer", "transform"})
+  .add_int64_power_of_two_axis("Elements{io}", nvbench::range(24, 28, 4));