Implement prefetching

cub/benchmarks/bench/reduce/by_key.cu

@@ -13,6 +13,7 @@
 // %RANGE% TUNE_MAGIC_NS ns 0:2048:4
 // %RANGE% TUNE_DELAY_CONSTRUCTOR_ID dcid 0:7:1
 // %RANGE% TUNE_L2_WRITE_LATENCY_NS l2w 0:1200:5
+// %RANGE% TUNE_LOAD_PREFETCH prf 0:3:1
 
 #if !TUNE_BASE
 struct bench_reduce_by_key_policy_selector
@@ -26,6 +27,7 @@ struct bench_reduce_by_key_policy_selector
       TUNE_LOAD == 0 ? cub::LOAD_DEFAULT : cub::LOAD_CA,
       cub::BLOCK_SCAN_WARP_SCANS,
       lookback_delay_policy,
+      static_cast<cub::detail::LoadPrefetch>(TUNE_LOAD_PREFETCH),
     };
   }
 };

cub/cub/agent/agent_reduce_by_key.cuh

@@ -23,6 +23,7 @@
 #include <cub/block/block_load.cuh>
 #include <cub/block/block_scan.cuh>
 #include <cub/block/block_store.cuh>
+#include <cub/detail/prefetch.cuh>
 #include <cub/iterator/cache_modified_input_iterator.cuh>
 
 #include <cuda/__functional/operator_properties.h>
@@ -43,14 +44,16 @@ template <int ThreadsPerBlock,
           BlockLoadAlgorithm LoadAlgorithm,
           CacheLoadModifier LoadModifier,
           BlockScanAlgorithm ScanAlgorithm,
-          typename DelayConstructorT = detail::fixed_delay_constructor_t<350, 450>>
+          typename DelayConstructorT             = detail::fixed_delay_constructor_t<350, 450>,
+          detail::LoadPrefetch LoadPrefetchLevel = detail::LoadPrefetch::none>
 struct agent_reduce_by_key_policy
 {
-  static constexpr int BLOCK_THREADS                 = ThreadsPerBlock;
-  static constexpr int ITEMS_PER_THREAD              = ItemsPerThread;
-  static constexpr BlockLoadAlgorithm LOAD_ALGORITHM = LoadAlgorithm;
-  static constexpr CacheLoadModifier LOAD_MODIFIER   = LoadModifier;
-  static constexpr BlockScanAlgorithm SCAN_ALGORITHM = ScanAlgorithm;
+  static constexpr int BLOCK_THREADS                  = ThreadsPerBlock;
+  static constexpr int ITEMS_PER_THREAD               = ItemsPerThread;
+  static constexpr BlockLoadAlgorithm LOAD_ALGORITHM  = LoadAlgorithm;
+  static constexpr CacheLoadModifier LOAD_MODIFIER    = LoadModifier;
+  static constexpr BlockScanAlgorithm SCAN_ALGORITHM  = ScanAlgorithm;
+  static constexpr detail::LoadPrefetch LOAD_PREFETCH = LoadPrefetchLevel;
 
   struct detail
   {
@@ -740,6 +743,34 @@ struct AgentReduceByKey
     // Remaining items (including this tile)
     OffsetT num_remaining = num_items - tile_offset;
 
+    // Fire prefetch hints before the decoupled look-back stall inside ConsumeTile, so the
+    // stall hides the L2 fetch latency and the tile's data has arrived by the time it is loaded.
+    // prefetch_tile is a no-op for LoadPrefetch::none, so no enablement check is needed here.
+    // The wrapped iterators apply the policy's LOAD_MODIFIER to the user's raw pointers; unlike a
+    // user-supplied CacheModifiedInputIterator, this internal wrapper expresses no intent to bypass
+    // caches, so unwrapping it for the prefetch hint is legitimate.
+    const int items_to_prefetch = static_cast<int>(::cuda::std::min(num_remaining, static_cast<OffsetT>(TILE_ITEMS)));
+    if constexpr (::cuda::std::is_pointer_v<KeysInputIteratorT>)
+    {
+      detail::prefetch_tile<BLOCK_THREADS, AgentReduceByKeyPolicyT::LOAD_PREFETCH>(
+        threadIdx.x, (d_keys_in + tile_offset).ptr, items_to_prefetch);
+    }
+    else
+    {
+      detail::prefetch_tile<BLOCK_THREADS, AgentReduceByKeyPolicyT::LOAD_PREFETCH>(
+        threadIdx.x, d_keys_in + tile_offset, items_to_prefetch);
+    }
+    if constexpr (::cuda::std::is_pointer_v<ValuesInputIteratorT>)
+    {
+      detail::prefetch_tile<BLOCK_THREADS, AgentReduceByKeyPolicyT::LOAD_PREFETCH>(
+        threadIdx.x, (d_values_in + tile_offset).ptr, items_to_prefetch);
+    }
+    else
+    {
+      detail::prefetch_tile<BLOCK_THREADS, AgentReduceByKeyPolicyT::LOAD_PREFETCH>(
+        threadIdx.x, d_values_in + tile_offset, items_to_prefetch);
+    }
+
     if (num_remaining > TILE_ITEMS)
     {
       // Not last tile

cub/cub/detail/prefetch.cuh

@@ -0,0 +1,136 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3
+
+//! @file
+//! Cooperative global-memory prefetch hints for tiles of data, placed explicitly by algorithm agents (or kernel
+//! authors) at points in the kernel schedule where the hint has lead time before the corresponding loads.
+
+#pragma once
+
+#include <cub/config.cuh>
+
+#if defined(_CCCL_IMPLICIT_SYSTEM_HEADER_GCC)
+#  pragma GCC system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_CLANG)
+#  pragma clang system_header
+#elif defined(_CCCL_IMPLICIT_SYSTEM_HEADER_MSVC)
+#  pragma system_header
+#endif // no system header
+
+#include <cub/util_type.cuh>
+
+#include <cuda/__cmath/round_up.h>
+#include <cuda/__memcpy_async/elect_one.h>
+#include <cuda/__memory/align_down.h>
+#include <cuda/__ptx/ptx_helper_functions.h>
+#include <cuda/std/__host_stdlib/ostream>
+#include <cuda/std/__iterator/concepts.h>
+#include <cuda/std/__memory/pointer_traits.h>
+
+#include <nv/target>
+
+CUB_NAMESPACE_BEGIN
+
+namespace detail
+{
+//! Enumerates the cache levels a tile prefetch hint can target.
+enum class LoadPrefetch : int
+{
+  //! No prefetch hint emitted.
+  none,
+  //! Emit an L2 prefetch hint (``prefetch.global.L2``) to all affected cache lines.
+  l2,
+  //! Emit an L1 prefetch hint (``prefetch.global.L1``) to all affected cache lines. Falls back to L2 on
+  //! architectures that do not support real L1 prefetch.
+  l1,
+  //! Emit a TMA bulk prefetch into L2 (``cp.async.bulk.prefetch``) for the whole tile.
+  //! Requires SM_90 or later (Hopper/Blackwell). Falls back to a no-op on older architectures.
+  bulk_l2,
+};
+
+#if _CCCL_HOSTED() && !defined(_CCCL_DOXYGEN_INVOKED)
+inline ::std::ostream& operator<<(::std::ostream& os, LoadPrefetch prefetch)
+{
+  switch (prefetch)
+  {
+    case LoadPrefetch::none:
+      return os << "detail::LoadPrefetch::none";
+    case LoadPrefetch::l2:
+      return os << "detail::LoadPrefetch::l2";
+    case LoadPrefetch::l1:
+      return os << "detail::LoadPrefetch::l1";
+    case LoadPrefetch::bulk_l2:
+      return os << "detail::LoadPrefetch::bulk_l2";
+    default:
+      return os << "<unknown detail::LoadPrefetch: " << static_cast<int>(prefetch) << ">";
+  }
+}
+#endif // _CCCL_HOSTED() && !_CCCL_DOXYGEN_INVOKED
+
+//! Prefetch hints require a raw address, so only contiguous iterators qualify. Notably,
+//! cub::CacheModifiedInputIterator does not: it routes loads through an explicit cache path
+//! (e.g. ``LOAD_NC``) that a prefetch hint would defeat.
+template <typename RandomAccessIterator>
+inline constexpr bool can_prefetch_from =
+  ::cuda::std::contiguous_iterator<RandomAccessIterator> && ::cuda::std::__can_to_address<RandomAccessIterator>;
+
+//! Cooperatively emit prefetch hints covering the tile ``[block_src_it, block_src_it + items_to_prefetch)``.
+//! For ``LoadPrefetch::l1``/``l2``, the calling thread block's threads stride the tile's address range in
+//! ``PrefetchStride``-byte steps, each issuing one hint per cache line. For ``LoadPrefetch::bulk_l2``, one
+//! elected thread issues a single TMA bulk prefetch for the whole tile. Silently a no-op for
+//! ``LoadPrefetch::none`` and for iterator types that do not satisfy ``can_prefetch_from``, so call sites
+//! need no enablement check of their own.
+template <int ThreadsPerBlock,
+          LoadPrefetch Prefetch = LoadPrefetch::l2,
+          int PrefetchStride    = 128,
+          typename RandomAccessIterator>
+_CCCL_DEVICE _CCCL_FORCEINLINE void
+prefetch_tile(int linear_tid, RandomAccessIterator block_src_it, int items_to_prefetch)
+{
+  if constexpr (Prefetch != LoadPrefetch::none && can_prefetch_from<RandomAccessIterator>)
+  {
+    using input_t                  = cub::detail::it_value_t<RandomAccessIterator>;
+    const unsigned int total_bytes = static_cast<unsigned int>(items_to_prefetch) * unsigned{sizeof(input_t)};
+    const auto* const src_ptr      = reinterpret_cast<const char*>(::cuda::std::to_address(block_src_it));
+
+    if constexpr (Prefetch == LoadPrefetch::bulk_l2)
+    {
+      // One elected thread issues a single TMA bulk prefetch for the whole tile.
+      // cp.async.bulk.prefetch is fire-and-forget: no commit_group/wait_group needed.
+      // Requires SM_90+; a no-op on older architectures.
+      NV_IF_TARGET(NV_PROVIDES_SM_90, (if (::cuda::device::__block_elect_one()) {
+                     // srcMem must be 16-byte aligned per PTX ISA; align base down and extend size to compensate
+                     const auto* const aligned_base  = ::cuda::align_down(src_ptr, 16);
+                     const unsigned int prefix       = static_cast<unsigned int>(src_ptr - aligned_base);
+                     const unsigned int aligned_size = ::cuda::round_up(total_bytes + prefix, 16u);
+                     if (aligned_size > 0)
+                     {
+                       asm volatile("cp.async.bulk.prefetch.L2.global [%0], %1;"
+                                    :
+                                    : "l"(::cuda::ptx::__as_ptr_gmem(aligned_base)), "r"(aligned_size)
+                                    : "memory");
+                     }
+                   }))
+    }
+    else
+    {
+      _CCCL_PRAGMA_NOUNROLL()
+      for (unsigned int offset = static_cast<unsigned int>(linear_tid) * PrefetchStride; offset < total_bytes;
+           offset += static_cast<unsigned int>(ThreadsPerBlock) * PrefetchStride)
+      {
+        // TODO: replace with cuda::ptx::prefetch_L1/L2 once exposed in libcudacxx
+        if constexpr (Prefetch == LoadPrefetch::l1)
+        {
+          asm volatile("prefetch.global.L1 [%0];" : : "l"(::cuda::ptx::__as_ptr_gmem(src_ptr + offset)) : "memory");
+        }
+        else
+        {
+          asm volatile("prefetch.global.L2 [%0];" : : "l"(::cuda::ptx::__as_ptr_gmem(src_ptr + offset)) : "memory");
+        }
+      }
+    }
+  }
+}
+} // namespace detail
+
+CUB_NAMESPACE_END

cub/cub/device/dispatch/dispatch_reduce_by_key.cuh

@@ -210,7 +210,8 @@ __launch_bounds__(int(current_policy<PolicySelector>().threads_per_block))
                 policy.load_algorithm,
                 policy.load_modifier,
                 policy.scan_algorithm,
-                delay_constructor_t<policy.lookback_delay.kind, policy.lookback_delay.delay, policy.lookback_delay.l2_write_latency>>;
+                delay_constructor_t<policy.lookback_delay.kind, policy.lookback_delay.delay, policy.lookback_delay.l2_write_latency>,
+                policy.load_prefetch>;
 
   using vsmem_helper_t = vsmem_helper_default_fallback_policy_t<
     AgentReduceByKeyPolicyT,
@@ -654,7 +655,8 @@ _CCCL_HOST_DEVICE_API auto determine_threads_items_vsmem(PolicyGetter policy_get
                           policy.load_algorithm,
                           policy.load_modifier,
                           policy.scan_algorithm,
-                          delay_constructor_t<policy.lookback_delay.kind, policy.lookback_delay.delay, policy.lookback_delay.l2_write_latency>>;
+                          delay_constructor_t<policy.lookback_delay.kind, policy.lookback_delay.delay, policy.lookback_delay.l2_write_latency>,
+                          policy.load_prefetch>;
   using vsmem_helper_t = vsmem_helper_default_fallback_policy_t<Policy, AgentReduceByKey, Args...>;
   return ::cuda::std::tuple{vsmem_helper_t::agent_policy_t::BLOCK_THREADS,
                             vsmem_helper_t::agent_policy_t::ITEMS_PER_THREAD,

cub/cub/device/dispatch/tuning/tuning_reduce_by_key.cuh

@@ -18,6 +18,7 @@
 #include <cub/block/block_load.cuh>
 #include <cub/block/block_scan.cuh>
 #include <cub/detail/delay_constructor.cuh>
+#include <cub/detail/prefetch.cuh>
 #include <cub/device/dispatch/tuning/common.cuh>
 #include <cub/util_device.cuh>
 #include <cub/util_type.cuh>
@@ -40,13 +41,16 @@ struct ReduceByKeyPolicy
   CacheLoadModifier load_modifier; //!< The @ref CacheLoadModifier used for loading items from global memory
   BlockScanAlgorithm scan_algorithm; //!< The @ref BlockScanAlgorithm used for the prefix scan
   LookbackDelayPolicy lookback_delay; //!< The @ref LookbackDelayPolicy used for the lookback delay
+  detail::LoadPrefetch load_prefetch = detail::LoadPrefetch::none; //!< The cache level the agent prefetches each
+                                                                   //!< input tile into ahead of its loads
 
   [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr friend bool
   operator==(const ReduceByKeyPolicy& lhs, const ReduceByKeyPolicy& rhs) noexcept
   {
     return lhs.threads_per_block == rhs.threads_per_block && lhs.items_per_thread == rhs.items_per_thread
         && lhs.load_algorithm == rhs.load_algorithm && lhs.load_modifier == rhs.load_modifier
-        && lhs.scan_algorithm == rhs.scan_algorithm && lhs.lookback_delay == rhs.lookback_delay;
+        && lhs.scan_algorithm == rhs.scan_algorithm && lhs.lookback_delay == rhs.lookback_delay
+        && lhs.load_prefetch == rhs.load_prefetch;
   }
 
   [[nodiscard]] _CCCL_HOST_DEVICE_API constexpr friend bool
@@ -59,9 +63,10 @@ struct ReduceByKeyPolicy
   friend ::std::ostream& operator<<(::std::ostream& os, const ReduceByKeyPolicy& p)
   {
     return os
-        << "ReduceByKeyPolicy { .threads_per_block = " << p.threads_per_block << ", .items_per_thread = "
-        << p.items_per_thread << ", .load_algorithm = " << p.load_algorithm << ", .load_modifier = " << p.load_modifier
-        << ", .scan_algorithm = " << p.scan_algorithm << ", .lookback_delay = " << p.lookback_delay << " }";
+        << "ReduceByKeyPolicy { .threads_per_block = " << p.threads_per_block
+        << ", .items_per_thread = " << p.items_per_thread << ", .load_algorithm = " << p.load_algorithm
+        << ", .load_modifier = " << p.load_modifier << ", .scan_algorithm = " << p.scan_algorithm
+        << ", .lookback_delay = " << p.lookback_delay << ", .load_prefetch = " << p.load_prefetch << " }";
   }
 #endif // _CCCL_HOSTED()
 };

cub/test/catch2_test_device_reduce_by_key_env.cu

@@ -0,0 +1,105 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+
+// Should precede any includes
+struct stream_registry_factory_t;
+#define CUB_DETAIL_DEFAULT_KERNEL_LAUNCHER_FACTORY stream_registry_factory_t
+
+#include "insert_nested_NVTX_range_guard.h"
+
+#include <cub/device/device_reduce.cuh>
+
+#include <thrust/device_vector.h>
+#include <thrust/host_vector.h>
+
+#include <cuda/devices>
+#include <cuda/std/functional>
+
+#include "catch2_test_env_launch_helper.h"
+
+DECLARE_LAUNCH_WRAPPER(cub::DeviceReduce::ReduceByKey, device_reduce_by_key);
+
+// %PARAM% TEST_LAUNCH lid 0:1:2
+
+#include <c2h/catch2_test_helper.h>
+
+template <cub::detail::LoadPrefetch PrefetchLevel>
+struct reduce_by_key_prefetch_tuning
+{
+  _CCCL_API constexpr auto operator()(cuda::compute_capability) const -> cub::ReduceByKeyPolicy
+  {
+    return {128, 11, cub::BLOCK_LOAD_DIRECT, cub::LOAD_DEFAULT, cub::BLOCK_SCAN_WARP_SCANS, {}, PrefetchLevel};
+  }
+};
+
+using load_prefetch_levels =
+  c2h::enum_type_list<cub::detail::LoadPrefetch,
+                      cub::detail::LoadPrefetch::none,
+                      cub::detail::LoadPrefetch::l2,
+                      cub::detail::LoadPrefetch::l1,
+                      cub::detail::LoadPrefetch::bulk_l2>;
+
+#if TEST_LAUNCH != 1
+
+// Prefetch hints must never affect results: run the same reduction under every LoadPrefetch level
+// (bulk_l2 is a no-op below SM90) and expect output identical to a sequential host reference.
+C2H_TEST("Device ReduceByKey is correct under every load_prefetch level", "[reduce][device]", load_prefetch_levels)
+{
+  constexpr auto prefetch_level = c2h::get<0, TestType>::value;
+
+  // Cover a single partial tile, and many tiles with a ragged final tile
+  const int num_items = GENERATE(values({42, (1 << 20) + 3}));
+  CAPTURE(num_items);
+
+  // Deterministic keys with short runs, and values that vary within each run
+  thrust::host_vector<int> h_keys(num_items);
+  thrust::host_vector<int> h_values(num_items);
+  for (int i = 0; i < num_items; ++i)
+  {
+    h_keys[i]   = (i / 3) % 5;
+    h_values[i] = i % 7;
+  }
+  thrust::device_vector<int> d_keys_in(h_keys);
+  thrust::device_vector<int> d_values_in(h_values);
+
+  c2h::device_vector<int> d_unique_out(num_items, thrust::no_init);
+  c2h::device_vector<int> d_aggregates_out(num_items, thrust::no_init);
+  c2h::device_vector<int> d_num_runs_out(1, thrust::no_init);
+
+  // Raw pointers on purpose: the agent wraps pointer inputs in CacheModifiedInputIterator and
+  // unwraps them for the prefetch hint — the code path production dispatch actually takes.
+  device_reduce_by_key(
+    thrust::raw_pointer_cast(d_keys_in.data()),
+    thrust::raw_pointer_cast(d_unique_out.data()),
+    thrust::raw_pointer_cast(d_values_in.data()),
+    thrust::raw_pointer_cast(d_aggregates_out.data()),
+    thrust::raw_pointer_cast(d_num_runs_out.data()),
+    cuda::std::plus<>{},
+    num_items,
+    cuda::execution::tune(reduce_by_key_prefetch_tuning<prefetch_level>{}));
+
+  // Sequential host reference
+  thrust::host_vector<int> expected_unique;
+  thrust::host_vector<int> expected_aggregates;
+  for (int i = 0; i < num_items; ++i)
+  {
+    if (i == 0 || h_keys[i] != h_keys[i - 1])
+    {
+      expected_unique.push_back(h_keys[i]);
+      expected_aggregates.push_back(h_values[i]);
+    }
+    else
+    {
+      expected_aggregates.back() += h_values[i];
+    }
+  }
+
+  const auto num_runs = static_cast<size_t>(d_num_runs_out[0]);
+  REQUIRE(num_runs == expected_unique.size());
+  d_unique_out.resize(num_runs);
+  d_aggregates_out.resize(num_runs);
+  REQUIRE(d_unique_out == expected_unique);
+  REQUIRE(d_aggregates_out == expected_aggregates);
+}
+
+#endif // TEST_LAUNCH != 1