Implement prefetching

cub/cub/detail/prefetch.cuh

@@ -0,0 +1,142 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+
+//! @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.
+// Note: ``cub::CacheModifiedInputIterator`` does not qualify: 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>;
+
+//! A block-scope collective that cooperatively emits global-memory prefetch hints for a tile. The element type
+//! ``T`` and block size ``ThreadsPerBlock`` are fixed on the type; the cache level (``PrefetchLevel``) and the
+//! per-cache-line hint spacing (``PrefetchStride``, bytes) are template configuration. Placement in the kernel
+//! schedule is the caller's — fire it where there is lead time before the corresponding loads.
+//!
+//! @tparam T             Element type of the tile being prefetched (drives the byte-size math).
+//! @tparam ThreadsPerBlock Number of threads in the block cooperating on the hints.
+//! @tparam PrefetchLevel  Which cache level to target (see ``LoadPrefetch``); ``none`` makes ``Prefetch()`` a no-op.
+//! @tparam PrefetchStride Byte stride between successive hints; one hint per cache line (default 128 B).
+template <typename T, int ThreadsPerBlock, LoadPrefetch PrefetchLevel = LoadPrefetch::l2, int PrefetchStride = 128>
+struct BlockPrefetch
+{
+  //! Cooperatively emit prefetch hints covering the tile ``[block_src_it, block_src_it + items_to_prefetch)``
+  template <typename RandomAccessIterator>
+  static _CCCL_DEVICE _CCCL_FORCEINLINE void Prefetch(RandomAccessIterator block_src_it, int items_to_prefetch)
+  {
+    if constexpr (PrefetchLevel != LoadPrefetch::none && can_prefetch_from<RandomAccessIterator>)
+    {
+      static_assert(sizeof(cub::detail::it_value_t<RandomAccessIterator>) == sizeof(T),
+                    "BlockPrefetch element type T must match the iterator's value type size");
+      const int linear_tid           = static_cast<int>(threadIdx.x);
+      const unsigned int total_bytes = static_cast<unsigned int>(items_to_prefetch) * unsigned{sizeof(T)};
+      const auto* const src_ptr      = reinterpret_cast<const char*>(::cuda::std::to_address(block_src_it));
+
+      if constexpr (PrefetchLevel == 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 (PrefetchLevel == 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/test/catch2_test_block_prefetch.cu

@@ -0,0 +1,138 @@
+// SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.
+// SPDX-License-Identifier: BSD-3-Clause
+
+#include <cub/detail/prefetch.cuh>
+#include <cub/iterator/cache_modified_input_iterator.cuh>
+
+#include <c2h/catch2_test_helper.h>
+
+// BlockPrefetch emits global-memory prefetch *hints*: they carry no functional result and cannot be observed from
+// the device program. A runtime test can therefore only assert that (1) issuing the hints never faults for any
+// level / element type / tile shape, (2) the hints never disturb the data they cover, and (3) `none` and
+// non-contiguous iterators (e.g. CacheModifiedInputIterator) compile and run as a genuine no-op. Whether the
+// intended SASS is actually emitted is verified out-of-band (cuobjdump), not here.
+
+// Compile-time coverage of the trait that gates every hint.
+static_assert(cub::detail::can_prefetch_from<int*>, "raw pointers are contiguous and must be prefetchable");
+static_assert(cub::detail::can_prefetch_from<const double*>, "const raw pointers must be prefetchable");
+static_assert(!cub::detail::can_prefetch_from<cub::CacheModifiedInputIterator<cub::CacheLoadModifier::LOAD_CS, int>>,
+              "CacheModifiedInputIterator routes through an explicit cache path and must be rejected");
+
+// Prefetch the tile, then faithfully copy it through, so the launch is observable and any corruption of the
+// prefetched region is caught by the host-side comparison.
+template <typename T, int ThreadsInBlock, cub::detail::LoadPrefetch Level, int Stride, typename InputIteratorT>
+__global__ void block_prefetch_kernel(InputIteratorT input, T* output, int num_items)
+{
+  cub::detail::BlockPrefetch<T, ThreadsInBlock, Level, Stride>::Prefetch(input, num_items);
+
+  for (int i = static_cast<int>(threadIdx.x); i < num_items; i += ThreadsInBlock)
+  {
+    output[i] = input[i];
+  }
+}
+
+template <typename T, int ThreadsInBlock, cub::detail::LoadPrefetch Level, int Stride = 128, typename InputIteratorT>
+void test_block_prefetch(const c2h::device_vector<T>& d_input, InputIteratorT input)
+{
+  const int num_items = static_cast<int>(d_input.size());
+  c2h::device_vector<T> d_output(num_items, T{});
+
+  block_prefetch_kernel<T, ThreadsInBlock, Level, Stride>
+    <<<1, ThreadsInBlock>>>(input, thrust::raw_pointer_cast(d_output.data()), num_items);
+  REQUIRE(cudaSuccess == cudaPeekAtLastError());
+  REQUIRE(cudaSuccess == cudaDeviceSynchronize());
+
+  // A prefetch hint must never alter the data it covers.
+  REQUIRE(d_input == d_output);
+}
+
+using types            = c2h::type_list<std::uint8_t, std::int32_t, std::int64_t>;
+using threads_in_block = c2h::enum_type_list<int, 32, 128>;
+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>;
+
+template <class TestType>
+struct params_t
+{
+  using type = typename c2h::get<0, TestType>;
+
+  static constexpr int threads_in_block            = c2h::get<1, TestType>::value;
+  static constexpr cub::detail::LoadPrefetch level = c2h::get<2, TestType>::value;
+};
+
+C2H_TEST("BlockPrefetch issues hints without disturbing the tile",
+         "[prefetch][block]",
+         types,
+         threads_in_block,
+         load_prefetch_levels)
+{
+  using params = params_t<TestType>;
+  using type   = typename params::type;
+
+  // Cover the empty tile, single item, sub-block and ragged tiles, and tiles spanning many cache lines.
+  constexpr int max_items = 8 * params::threads_in_block;
+  const int num_items     = GENERATE_COPY(0, 1, 7, params::threads_in_block + 3, take(5, random(1, max_items)));
+  CAPTURE(num_items);
+
+  c2h::device_vector<type> d_input(num_items);
+  c2h::gen(C2H_SEED(2), d_input);
+
+  test_block_prefetch<type, params::threads_in_block, params::level>(d_input, thrust::raw_pointer_cast(d_input.data()));
+}
+
+C2H_TEST("BlockPrefetch is a no-op for CacheModifiedInputIterator", "[prefetch][block]", load_prefetch_levels)
+{
+  using type                                = int;
+  constexpr int threads_in_block            = 128;
+  constexpr cub::detail::LoadPrefetch level = c2h::get<0, TestType>::value;
+
+  const int num_items = GENERATE_COPY(7, 200, take(3, random(1, 1024)));
+  CAPTURE(num_items);
+
+  c2h::device_vector<type> d_input(num_items);
+  c2h::gen(C2H_SEED(2), d_input);
+
+  // can_prefetch_from<CMI> is false, so Prefetch must compile out to nothing; the copy still reads through the CMI.
+  cub::CacheModifiedInputIterator<cub::CacheLoadModifier::LOAD_CS, type> in(thrust::raw_pointer_cast(d_input.data()));
+  test_block_prefetch<type, threads_in_block, level>(d_input, in);
+}
+
+C2H_TEST("BlockPrefetch handles unaligned tile bases", "[prefetch][block]", c2h::type_list<std::uint8_t, std::int32_t>)
+{
+  using type                     = c2h::get<0, TestType>;
+  constexpr int threads_in_block = 128;
+
+  // bulk_l2 aligns the base down to 16 B and extends the size to compensate; walk the base across a 16 B window.
+  const int offset    = GENERATE(0, 1, 2, 3, 4, 5, 7, 8, 15);
+  const int num_items = GENERATE(1, 33, 512);
+  CAPTURE(offset, num_items);
+
+  c2h::device_vector<type> d_storage(num_items + offset);
+  c2h::gen(C2H_SEED(1), d_storage);
+
+  // Prefetch/copy the sub-range that starts at an unaligned offset into the allocation.
+  auto* base = thrust::raw_pointer_cast(d_storage.data()) + offset;
+  c2h::device_vector<type> d_input(d_storage.begin() + offset, d_storage.end());
+
+  test_block_prefetch<type, threads_in_block, cub::detail::LoadPrefetch::bulk_l2>(d_input, base);
+}
+
+C2H_TEST("BlockPrefetch honors a non-default stride", "[prefetch][block]")
+{
+  using type                     = int;
+  constexpr int threads_in_block = 64;
+
+  const int num_items = GENERATE_COPY(values({1, 100, 777}));
+  CAPTURE(num_items);
+
+  c2h::device_vector<type> d_input(num_items);
+  c2h::gen(C2H_SEED(1), d_input);
+
+  // A coarser 256 B stride issues fewer hints; the data must still be left intact.
+  test_block_prefetch<type, threads_in_block, cub::detail::LoadPrefetch::l2, 256>(
+    d_input, thrust::raw_pointer_cast(d_input.data()));
+}