Add prefetch transform kernel

Fixes: #2363
NVIDIA · Sep 10, 2024 · 0cfe4ba · 0cfe4ba
1 parent ee9b856
commit 0cfe4ba
Show file tree

Hide file tree

Showing 6 changed files with 136 additions and 15 deletions.
diff --git a/cub/benchmarks/bench/transform/babelstream.h b/cub/benchmarks/bench/transform/babelstream.h
@@ -23,7 +23,9 @@ struct policy_hub_t
     using algo_policy =
       ::cuda::std::_If<algorithm == cub::detail::transform::Algorithm::fallback_for,
                        cub::detail::transform::fallback_for_policy,
-                       cub::detail::transform::async_copy_policy_t<TUNE_THREADS>>;
+                       ::cuda::std::_If<algorithm == cub::detail::transform::Algorithm::prefetch,
+                                        cub::detail::transform::prefetch_policy_t<TUNE_THREADS>,
+                                        cub::detail::transform::async_copy_policy_t<TUNE_THREADS>>>;
   };
 };
 #endif

diff --git a/cub/benchmarks/bench/transform/babelstream1.cu b/cub/benchmarks/bench/transform/babelstream1.cu
@@ -2,15 +2,15 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
 // %RANGE% TUNE_THREADS tpb 128:1024:128
-// %RANGE% TUNE_ALGORITHM alg 0:1:1
+// %RANGE% TUNE_ALGORITHM alg 0:2:1
 
 // keep checks at the top so compilation of discarded variants fails really fast
 #if !TUNE_BASE
-#  if TUNE_ALGORITHM == 1 && (__CUDA_ARCH_LIST__) < 900
+#  if TUNE_ALGORITHM == 2 && (__CUDA_ARCH_LIST__) < 900
 #    error "Cannot compile algorithm 4 (ublkcp) below sm90"
 #  endif
 
-#  if TUNE_ALGORITHM == 1 && !defined(_CUB_HAS_TRANSFORM_UBLKCP)
+#  if TUNE_ALGORITHM == 2 && !defined(_CUB_HAS_TRANSFORM_UBLKCP)
 #    error "Cannot tune for ublkcp algorithm, which is not provided by CUB (old CTK?)"
 #  endif
 #endif

diff --git a/cub/benchmarks/bench/transform/babelstream2.cu b/cub/benchmarks/bench/transform/babelstream2.cu
@@ -2,15 +2,15 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
 // %RANGE% TUNE_THREADS tpb 128:1024:128
-// %RANGE% TUNE_ALGORITHM alg 0:1:1
+// %RANGE% TUNE_ALGORITHM alg 0:2:1
 
 // keep checks at the top so compilation of discarded variants fails really fast
 #if !TUNE_BASE
-#  if TUNE_ALGORITHM == 1 && (__CUDA_ARCH_LIST__) < 900
+#  if TUNE_ALGORITHM == 2 && (__CUDA_ARCH_LIST__) < 900
 #    error "Cannot compile algorithm 4 (ublkcp) below sm90"
 #  endif
 
-#  if TUNE_ALGORITHM == 1 && !defined(_CUB_HAS_TRANSFORM_UBLKCP)
+#  if TUNE_ALGORITHM == 2 && !defined(_CUB_HAS_TRANSFORM_UBLKCP)
 #    error "Cannot tune for ublkcp algorithm, which is not provided by CUB (old CTK?)"
 #  endif
 #endif

diff --git a/cub/benchmarks/bench/transform/babelstream3.cu b/cub/benchmarks/bench/transform/babelstream3.cu
@@ -2,15 +2,15 @@
 // SPDX-License-Identifier: BSD-3-Clause
 
 // %RANGE% TUNE_THREADS tpb 128:1024:128
-// %RANGE% TUNE_ALGORITHM alg 0:1:1
+// %RANGE% TUNE_ALGORITHM alg 0:2:1
 
 // keep checks at the top so compilation of discarded variants fails really fast
 #if !TUNE_BASE
-#  if TUNE_ALGORITHM == 1 && (__CUDA_ARCH_LIST__) < 900
+#  if TUNE_ALGORITHM == 2 && (__CUDA_ARCH_LIST__) < 900
 #    error "Cannot compile algorithm 4 (ublkcp) below sm90"
 #  endif
 
-#  if TUNE_ALGORITHM == 1 && !defined(_CUB_HAS_TRANSFORM_UBLKCP)
+#  if TUNE_ALGORITHM == 2 && !defined(_CUB_HAS_TRANSFORM_UBLKCP)
 #    error "Cannot tune for ublkcp algorithm, which is not provided by CUB (old CTK?)"
 #  endif
 #endif

diff --git a/cub/cub/device/dispatch/dispatch_transform.cuh b/cub/cub/device/dispatch/dispatch_transform.cuh
@@ -89,6 +89,7 @@ _CCCL_HOST_DEVICE constexpr auto loaded_bytes_per_iteration() -> int
 enum class Algorithm
 {
   fallback_for,
+  prefetch,
 #ifdef _CUB_HAS_TRANSFORM_UBLKCP
   ublkcp,
 #endif // _CUB_HAS_TRANSFORM_UBLKCP
@@ -133,6 +134,90 @@ _CCCL_DEVICE void transform_kernel_impl(
   }
 }
 
+template <int BlockThreads>
+struct prefetch_policy_t
+{
+  static constexpr int block_threads = BlockThreads;
+  // items per tile are determined at runtime. these (inclusive) bounds allow overriding that value via a tuning policy
+  static constexpr int items_per_thread_no_input = 2; // when there are no input iterators, the kernel is just filling
+  static constexpr int min_items_per_thread      = 1;
+  static constexpr int max_items_per_thread      = 32;
+};
+
+// Prefetches (at least on Hopper) a 128 byte cache line. Prefetching out-of-bounds addresses has no side effects
+// TODO(bgruber): there is also the cp.async.bulk.prefetch instruction available on Hopper. May improve perf a tiny bit
+// as we need to create less instructions to prefetch the same amount of data.
+template <typename T>
+_CCCL_DEVICE _CCCL_FORCEINLINE void prefetch(const T* addr)
+{
+  assert(__isGlobal(addr));
+  // TODO(bgruber): prefetch to L1 may be even better
+  asm volatile("prefetch.global.L2 [%0];" : : "l"(addr) : "memory");
+}
+
+// overload for any iterator that is not a pointer, do nothing
+template <typename It, ::cuda::std::__enable_if_t<!::cuda::std::is_pointer<It>::value, int> = 0>
+_CCCL_DEVICE _CCCL_FORCEINLINE void prefetch(It)
+{}
+
+// this kernel guarantees stable addresses for the parameters of the user provided function
+template <typename PrefetchPolicy,
+          typename Offset,
+          typename F,
+          typename RandomAccessIteratorOut,
+          typename... RandomAccessIteratorIn>
+_CCCL_DEVICE void transform_kernel_impl(
+  ::cuda::std::integral_constant<Algorithm, Algorithm::prefetch>,
+  Offset num_items,
+  int num_elem_per_thread,
+  F f,
+  RandomAccessIteratorOut out,
+  RandomAccessIteratorIn... ins)
+{
+  constexpr int block_dim = PrefetchPolicy::block_threads;
+  const int tile_stride   = block_dim * num_elem_per_thread;
+  const Offset offset     = static_cast<Offset>(blockIdx.x) * tile_stride;
+  const int tile_size     = static_cast<int>(::cuda::std::min(num_items - offset, Offset{tile_stride}));
+
+  // move index and iterator domain to the block/thread index, to reduce arithmetic in the loops below
+  {
+    int dummy[] = {(ins += offset, 0)..., 0};
+    (void) &dummy;
+    out += offset;
+  }
+
+  for (int j = 0; j < num_elem_per_thread; ++j)
+  {
+    const int idx = j * block_dim + threadIdx.x;
+    // TODO(bgruber): replace by fold over comma in C++17
+    int dummy[] = {(prefetch(ins + idx), 0)..., 0}; // extra zero to handle empty packs
+    (void) &dummy; // nvcc 11.1 needs extra strong unused warning suppression
+  }
+
+#define PREFETCH_AGENT(full_tile)                                                                                  \
+  /* ahendriksen: various unrolling yields less <1% gains at much higher compile-time cost */                      \
+  /* TODO(bgruber): A6000 disagrees */                                                                             \
+  _Pragma("unroll 1") for (int j = 0; j < num_elem_per_thread; ++j)                                                \
+  {                                                                                                                \
+    const int idx = j * block_dim + threadIdx.x;                                                                   \
+    if (full_tile || idx < tile_size)                                                                              \
+    {                                                                                                              \
+      /* we have to unwrap Thrust's proxy references here for backward compatibility (try zip_iterator.cu test) */ \
+      out[idx] = f(THRUST_NS_QUALIFIER::raw_reference_cast(ins[idx])...);                                          \
+    }                                                                                                              \
+  }
+
+  if (tile_stride == tile_size)
+  {
+    PREFETCH_AGENT(true);
+  }
+  else
+  {
+    PREFETCH_AGENT(false);
+  }
+#undef PREFETCH_AGENT
+}
+
 template <int BlockThreads>
 struct async_copy_policy_t
 {
@@ -543,8 +628,8 @@ struct policy_hub<RequiresStableAddress, ::cuda::std::tuple<RandomAccessIterator
   {
     static constexpr int min_bif = arch_to_min_bytes_in_flight(300);
     // TODO(bgruber): we don't need algo, because we can just detect the type of algo_policy
-    static constexpr auto algorithm = Algorithm::fallback_for;
-    using algo_policy               = fallback_for_policy;
+    static constexpr auto algorithm = Algorithm::prefetch;
+    using algo_policy               = prefetch_policy_t<256>;
   };
 
 #ifdef _CUB_HAS_TRANSFORM_UBLKCP
@@ -566,8 +651,8 @@ struct policy_hub<RequiresStableAddress, ::cuda::std::tuple<RandomAccessIterator
 
     static constexpr bool use_fallback =
       RequiresStableAddress || !can_memcpy || no_input_streams || exhaust_smem || any_type_is_overalinged;
-    static constexpr auto algorithm = use_fallback ? Algorithm::fallback_for : Algorithm::ublkcp;
-    using algo_policy               = ::cuda::std::_If<use_fallback, fallback_for_policy, async_policy>;
+    static constexpr auto algorithm = use_fallback ? Algorithm::prefetch : Algorithm::ublkcp;
+    using algo_policy               = ::cuda::std::_If<use_fallback, prefetch_policy_t<256>, async_policy>;
   };
 
   using max_policy = policy900;
@@ -828,6 +913,38 @@ struct dispatch_t<RequiresStableAddress,
           make_iterator_kernel_arg(THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(::cuda::std::get<Is>(in)))...));
   }
 
+  template <typename ActivePolicy, std::size_t... Is>
+  CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t
+  invoke_algorithm(cuda::std::index_sequence<Is...>, ::cuda::std::integral_constant<Algorithm, Algorithm::prefetch>)
+  {
+    using policy_t          = typename ActivePolicy::algo_policy;
+    constexpr int block_dim = policy_t::block_threads;
+    int max_occupancy       = 0;
+    const auto error        = CubDebug(MaxSmOccupancy(max_occupancy, CUB_DETAIL_TRANSFORM_KERNEL_PTR, block_dim, 0));
+    if (error != cudaSuccess)
+    {
+      return error;
+    }
+
+    const int items_per_thread =
+      loaded_bytes_per_iter == 0
+        ? +policy_t::items_per_thread_no_input
+        : ::cuda::ceil_div(ActivePolicy::min_bif, max_occupancy * block_dim * loaded_bytes_per_iter);
+    const int items_per_thread_clamped =
+      ::cuda::std::clamp(items_per_thread, +policy_t::min_items_per_thread, +policy_t::max_items_per_thread);
+    const int tile_size = block_dim * items_per_thread_clamped;
+    const auto grid_dim = static_cast<unsigned int>(::cuda::ceil_div(num_items, Offset{tile_size}));
+    return CubDebug(
+      THRUST_NS_QUALIFIER::cuda_cub::launcher::triple_chevron(grid_dim, block_dim, 0, stream)
+        .doit(
+          CUB_DETAIL_TRANSFORM_KERNEL_PTR,
+          num_items,
+          items_per_thread_clamped,
+          op,
+          out,
+          make_iterator_kernel_arg(THRUST_NS_QUALIFIER::try_unwrap_contiguous_iterator(::cuda::std::get<Is>(in)))...));
+  }
+
   template <typename ActivePolicy>
   CUB_RUNTIME_FUNCTION _CCCL_FORCEINLINE cudaError_t Invoke()
   {

diff --git a/cub/test/catch2_test_device_transform.cu b/cub/test/catch2_test_device_transform.cu
@@ -34,7 +34,9 @@ struct policy_hub_for_alg
     using algo_policy =
       ::cuda::std::_If<Alg == Algorithm::fallback_for,
                        cub::detail::transform::fallback_for_policy,
-                       cub::detail::transform::async_copy_policy_t<256>>;
+                       ::cuda::std::_If<Alg == Algorithm::prefetch,
+                                        cub::detail::transform::prefetch_policy_t<256>,
+                                        cub::detail::transform::async_copy_policy_t<256>>>;
   };
 };