[cudax] Harden cooperative Reduce reuse

cudax/include/cuda/experimental/__coop/reduce.cuh

@@ -23,6 +23,7 @@
 
 #include <cub/block/block_reduce.cuh>
 #include <cub/thread/thread_reduce.cuh>
+#include <cub/util_type.cuh>
 #include <cub/warp/warp_reduce.cuh>
 
 #include <cuda/__cmath/ceil_div.h>
@@ -37,6 +38,7 @@
 #include <cuda/std/optional>
 
 #include <cuda/experimental/__coop/shuffle_down.cuh>
+#include <cuda/experimental/__group/queries.cuh>
 #include <cuda/experimental/__utility/result_policy.cuh>
 #include <cuda/experimental/group.cuh>
 
@@ -123,15 +125,20 @@ template <bool _Broadcasted, class _Hierarchy, class _Tp, cuda::std::size_t _Np,
   const auto __result = _BlockReduce{__scratch.__block_reduce_}.Reduce(__thread_data, __red_fn);
   if constexpr (_Broadcasted)
   {
+    // Wait until all threads are done using the aliased CUB storage.
+    __group.sync_aligned();
     if (gpu_thread.is_root_rank(__group))
     {
       __scratch.__bcast_ = __result;
     }
     __group.sync_aligned();
-    return __scratch.__bcast_;
+    const auto __broadcast_result = __scratch.__bcast_;
+    __group.sync_aligned();
+    return __broadcast_result;
   }
   else
   {
+    __group.sync_aligned();
     return (gpu_thread.is_root_rank(__group)) ? ::cuda::std::optional{__result} : ::cuda::std::nullopt;
   }
 }
@@ -221,6 +228,7 @@ template <bool _Broadcasted, class _Hierarchy, class _Tp, cuda::std::size_t _Np,
     }
     else
     {
+      __group.sync_aligned();
       return (gpu_thread.is_root_rank(__group)) ? ::cuda::std::optional{__result} : ::cuda::std::nullopt;
     }
   }
@@ -287,15 +295,16 @@ template <bool _Broadcasted, class _Hierarchy, class _Tp, cuda::std::size_t _Np,
   }
   else
   {
+    __group.sync_aligned();
     return __result;
   }
 }
 
 _CCCL_TEMPLATE(bool _Broadcasted, class _Group, class _Tp, ::cuda::std::size_t _Np, class _RedFn)
 _CCCL_REQUIRES(::cuda::std::is_same_v<thread_level, typename _Group::unit_type>
                  _CCCL_AND ::cuda::std::is_same_v<warp_level, typename _Group::level_type>)
-[[nodiscard]] _CCCL_DEVICE_API auto
-__reduce_impl(::cuda::std::bool_constant<_Broadcasted>, _Group __group, _Tp (&__thread_data)[_Np], _RedFn __red_fn)
+[[nodiscard]] _CCCL_DEVICE_API auto __reduce_impl(
+  ::cuda::std::bool_constant<_Broadcasted>, const _Group& __group, _Tp (&__thread_data)[_Np], _RedFn __red_fn)
 {
   using _MappingResult         = typename _Group::__mapping_result_type;
   const auto& __mapping_result = __group.__mapping_result();
@@ -328,66 +337,88 @@ __reduce_impl(::cuda::std::bool_constant<_Broadcasted>, _Group __group, _Tp (&__
 _CCCL_TEMPLATE(bool _Broadcasted, class _Group, class _Tp, ::cuda::std::size_t _Np, class _RedFn)
 _CCCL_REQUIRES(::cuda::std::is_same_v<warp_level, typename _Group::unit_type>
                  _CCCL_AND ::cuda::std::is_same_v<block_level, typename _Group::level_type>)
-[[nodiscard]] _CCCL_DEVICE_API auto
-__reduce_impl(::cuda::std::bool_constant<_Broadcasted>, _Group __group, _Tp (&__thread_data)[_Np], _RedFn __red_fn)
+[[nodiscard]] _CCCL_DEVICE_API auto __reduce_impl(
+  ::cuda::std::bool_constant<_Broadcasted>, const _Group& __group, _Tp (&__thread_data)[_Np], _RedFn __red_fn)
 {
-  constexpr auto __nwarps_in_group = warp.static_count(__group);
+  using _MappingResult = typename _Group::__mapping_result_type;
+
+  constexpr auto __nwarps_in_group = ::cuda::experimental::__static_count_query_group<warp_level, _Group>();
   static_assert(__nwarps_in_group != ::cuda::std::dynamic_extent,
                 "cuda::coop::reduce requires the group to have statically known size");
+  static_assert(_MappingResult::is_always_contiguous(),
+                "cuda::coop::reduce requires the group mapping to be contiguous");
 
-  using _WarpReduce = ::cub::WarpReduce<_Tp>;
-  struct _AdditionalScratch
-  {
-    _Tp __partials_[__nwarps_in_group];
-    _Tp __bcast_;
-  };
+  using _BlockExts = decltype(gpu_thread.extents(block, __group.hierarchy()));
 
-  union _Scratch
+  // CUDA thread blocks contain at most 1024 threads (32 warps). Prefer a smaller
+  // bound when all block extents are known at compile time.
+  constexpr auto __nthreads_in_block =
+    (_BlockExts::rank_dynamic() == 0)
+      ? _BlockExts::static_extent(0) * _BlockExts::static_extent(1) * _BlockExts::static_extent(2)
+      : ::cuda::std::size_t{1024};
+  constexpr auto __nwarps_in_block = ::cuda::ceil_div(__nthreads_in_block, ::cuda::std::size_t{32});
+
+  using _WarpReduce = ::cub::WarpReduce<_Tp>;
+  struct _Scratch
   {
-    typename _WarpReduce::TempStorage __warp_reduce_[__nwarps_in_group];
-    _AdditionalScratch __additional_;
+    typename _WarpReduce::TempStorage __warp_reduce_[__nwarps_in_block];
+    alignas(_Tp)::cub::Uninitialized<_Tp[__nwarps_in_block]> __values_;
   };
   __shared__ _Scratch __scratch;
 
-  const auto __partial = _WarpReduce{__scratch.__warp_reduce_[warp.rank(__group)]}.Reduce(__thread_data, __red_fn);
-  __group.sync_aligned();
-
   this_warp __warp{__group.hierarchy()};
+  const auto __warp_rank_in_block = __warp.rank(block);
+  const auto __warp_rank_in_group = warp.rank(__group);
+  // This relies on contiguous mappings assigning unit ranks in ascending
+  // physical-warp order, so every warp derives the same first physical rank.
+  const auto __first_warp_rank = __warp_rank_in_block - __warp_rank_in_group;
+  _CCCL_ASSERT(__warp_rank_in_block < __nwarps_in_block, "invalid physical warp rank");
+  _CCCL_ASSERT(__first_warp_rank + __nwarps_in_group <= __nwarps_in_block, "invalid physical warp range");
+
+  // Physical-warp slots remain unique when independent mappings reuse group
+  // ranks or overlap in time.
+  auto& __values = __scratch.__values_.Alias();
+
+  const auto __partial = _WarpReduce{__scratch.__warp_reduce_[__warp_rank_in_block]}.Reduce(__thread_data, __red_fn);
+
   if (gpu_thread.is_root_rank(__warp))
   {
-    __scratch.__additional_.__partials_[warp.rank(__group)] = __partial;
+    __values[__warp_rank_in_block] = __partial;
   }
   __group.sync_aligned();
 
   _Tp __result;
   if (warp.is_root_rank(__group))
   {
     const auto __value = (gpu_thread.rank(__warp) < __nwarps_in_group)
-                         ? __scratch.__additional_.__partials_[gpu_thread.rank(__warp)]
+                         ? __values[__first_warp_rank + gpu_thread.rank(__warp)]
                          : ::cuda::identity_element<_RedFn, _Tp>();
-    __result           = _WarpReduce{__scratch.__warp_reduce_[0]}.Reduce(__value, __red_fn);
+    __result           = _WarpReduce{__scratch.__warp_reduce_[__first_warp_rank]}.Reduce(__value, __red_fn);
   }
 
   if constexpr (_Broadcasted)
   {
     if (gpu_thread.is_root_rank(__group))
     {
-      __scratch.__additional_.__bcast_ = __result;
+      __values[__first_warp_rank] = __result;
     }
     __group.sync_aligned();
-    return __scratch.__additional_.__bcast_;
+    const auto __broadcast_result = __values[__first_warp_rank];
+    __group.sync_aligned();
+    return __broadcast_result;
   }
   else
   {
+    __group.sync_aligned();
     return (gpu_thread.is_root_rank(__group)) ? ::cuda::std::optional{__result} : ::cuda::std::nullopt;
   }
 }
 
 template <class _Group, class _Tp, ::cuda::std::size_t _Np, class _RedFn>
 [[nodiscard]] _CCCL_DEVICE_API ::cuda::std::optional<_Tp>
-reduce(_Group __group, _Tp (&__thread_data)[_Np], _RedFn __red_fn)
+reduce(const _Group& __group, _Tp (&__thread_data)[_Np], _RedFn __red_fn)
 {
-  static_assert(gpu_thread.static_count(__group) != ::cuda::std::dynamic_extent,
+  static_assert(::cuda::experimental::__static_count_query_group<thread_level, _Group>() != ::cuda::std::dynamic_extent,
                 "cuda::coop::reduce requires the group to have statically known size");
 
   _CCCL_ASSERT(gpu_thread.is_part_of(__group), "Only threads that are part of the group can call cudax::coop::reduce");
@@ -396,9 +427,10 @@ reduce(_Group __group, _Tp (&__thread_data)[_Np], _RedFn __red_fn)
 }
 
 template <class _Group, class _Tp, ::cuda::std::size_t _Np, class _RedFn>
-[[nodiscard]] _CCCL_DEVICE_API _Tp reduce(broadcasted_t, _Group __group, _Tp (&__thread_data)[_Np], _RedFn __red_fn)
+[[nodiscard]] _CCCL_DEVICE_API _Tp
+reduce(broadcasted_t, const _Group& __group, _Tp (&__thread_data)[_Np], _RedFn __red_fn)
 {
-  static_assert(gpu_thread.static_count(__group) != ::cuda::std::dynamic_extent,
+  static_assert(::cuda::experimental::__static_count_query_group<thread_level, _Group>() != ::cuda::std::dynamic_extent,
                 "cuda::coop::reduce requires the group to have statically known size");
 
   _CCCL_ASSERT(gpu_thread.is_part_of(__group), "Only threads that are part of the group can call cudax::coop::reduce");

cudax/include/cuda/experimental/__group/fwd.cuh

@@ -21,9 +21,13 @@
 #  pragma system_header
 #endif // no system header
 
-// Q: Do we want to enable this by default, or do we want the user to define some macro to get the interoperability with
-//    cooperative groups?
-#if __has_include(<cooperative_groups.h>)
+// Internal hook for embedded compilation environments that cannot consume
+// cooperative_groups.h. Define it consistently for every translation unit in
+// a program, before any cudax group header, to suppress cooperative-groups
+// includes and conversion overloads.
+#if defined(_CUDAX_DISABLE_COOPERATIVE_GROUPS_INTEROP)
+#  define _CCCL_HAS_COOPERATIVE_GROUPS() 0
+#elif __has_include(<cooperative_groups.h>)
 #  define _CCCL_HAS_COOPERATIVE_GROUPS() 1
 #else // ^^^ has cooperative groups ^^^ / vvv no cooperative groups vvv
 #  define _CCCL_HAS_COOPERATIVE_GROUPS() 0

cudax/test/CMakeLists.txt

@@ -183,6 +183,9 @@ cudax_add_catch2_test(test_target group.implicit_hierarchy
 cudax_add_catch2_test(test_target group.segmented_algorithm
     group/segmented_algorithm.cu
 )
+cudax_add_catch2_test(test_target group.no_cooperative_groups_interop
+    group/no_cooperative_groups_interop.cu
+)
 cudax_add_catch2_test(test_target group.this_group
     group/this_group.cu
 )

cudax/test/coop/reduce/this_block.cu

@@ -26,6 +26,9 @@
 #include <c2h/generators.h>
 #include <catch2/matchers/catch_matchers_floating_point.hpp>
 
+inline constexpr int reuse_block_size = 128;
+inline constexpr int reuse_repeats    = 16;
+
 /***********************************************************************************************************************
  * Thread Reduce Wrapper Kernels
  **********************************************************************************************************************/
@@ -71,6 +74,37 @@ struct ReduceKernel
   }
 };
 
+template <bool Broadcasted>
+struct RepeatedReduceKernel
+{
+  template <class Config>
+  __device__ void operator()(Config config, int* d_out)
+  {
+    cudax::this_block block{config};
+    const int thread_rank = cuda::gpu_thread.rank_as<int>(block);
+    constexpr int stride  = Broadcasted ? reuse_block_size : 1;
+
+    for (int repeat = 0; repeat < reuse_repeats; ++repeat)
+    {
+      int thread_data[1] = {repeat + 1};
+      if constexpr (Broadcasted)
+      {
+        const auto result = cudax::coop::reduce(cudax::broadcasted, block, thread_data, cuda::std::plus<>{});
+        d_out[repeat * stride + thread_rank] = result;
+      }
+      else
+      {
+        const auto result = cudax::coop::reduce(block, thread_data, cuda::std::plus<>{});
+        REQUIRE(result.has_value() == cuda::gpu_thread.is_root_rank(block));
+        if (result.has_value())
+        {
+          d_out[repeat] = *result;
+        }
+      }
+    }
+  }
+};
+
 /***********************************************************************************************************************
  * Type list definition
  **********************************************************************************************************************/
@@ -139,8 +173,37 @@ void run_reduce_kernel(
   stream.sync();
 }
 
-constexpr int max_size  = 4;
-constexpr int num_seeds = 10;
+template <bool Broadcasted>
+void run_repeated_reduce_kernel(cuda::stream_ref stream)
+{
+  constexpr int stride      = Broadcasted ? reuse_block_size : 1;
+  constexpr int output_size = reuse_repeats * stride;
+  c2h::device_vector<int> d_out(output_size, -1);
+  const auto out_ptr = thrust::raw_pointer_cast(d_out.data());
+  const auto config  = cuda::make_config(cuda::grid_dims<1>(), cuda::block_dims<reuse_block_size>());
+  cuda::launch(stream, config, RepeatedReduceKernel<Broadcasted>{}, out_ptr);
+  stream.sync();
+
+  const c2h::host_vector<int> h_out = d_out;
+  for (int repeat = 0; repeat < reuse_repeats; ++repeat)
+  {
+    const int expected = reuse_block_size * (repeat + 1);
+    if constexpr (Broadcasted)
+    {
+      for (int rank = 0; rank < reuse_block_size; ++rank)
+      {
+        REQUIRE(h_out[repeat * stride + rank] == expected);
+      }
+    }
+    else
+    {
+      REQUIRE(h_out[repeat] == expected);
+    }
+  }
+}
+
+inline constexpr int max_size  = 4;
+inline constexpr int num_seeds = 10;
 
 /***********************************************************************************************************************
  * Test cases
@@ -221,3 +284,11 @@ C2H_TEST("reduce/this_block Broadcasted", "[reduce][this_block]", integral_type_
     verify_results(c2h::host_vector<value_t>(block_size_t::value, reference_result), c2h::host_vector<value_t>(d_out));
   }
 }
+
+C2H_TEST("reduce/this_block can reuse scratch", "[reduce][this_block]")
+{
+  cuda::stream stream{cuda::devices[0]};
+
+  run_repeated_reduce_kernel<false>(stream);
+  run_repeated_reduce_kernel<true>(stream);
+}

cudax/test/coop/reduce/this_cluster.cu

@@ -26,7 +26,9 @@
 #include <c2h/generators.h>
 #include <catch2/matchers/catch_matchers_floating_point.hpp>
 
-constexpr int block_size = 64;
+inline constexpr int block_size         = 64;
+inline constexpr int reuse_cluster_size = 2;
+inline constexpr int reuse_repeats      = 16;
 
 /***********************************************************************************************************************
  * Thread Reduce Wrapper Kernels
@@ -70,6 +72,26 @@ struct ReduceKernel
   }
 };
 
+struct RepeatedReduceKernel
+{
+  template <class Config>
+  __device__ void operator()(Config config, int* d_out)
+  {
+    cudax::this_cluster cluster{config};
+
+    for (int repeat = 0; repeat < reuse_repeats; ++repeat)
+    {
+      int thread_data[1] = {repeat + 1};
+      const auto result  = cudax::coop::reduce(cluster, thread_data, cuda::std::plus<>{});
+      REQUIRE(result.has_value() == cuda::gpu_thread.is_root_rank(cluster));
+      if (result.has_value())
+      {
+        d_out[repeat] = *result;
+      }
+    }
+  }
+};
+
 /***********************************************************************************************************************
  * Type list definition
  **********************************************************************************************************************/
@@ -139,8 +161,25 @@ void run_reduce_kernel(
   stream.sync();
 }
 
-constexpr int max_size  = 4;
-constexpr int num_seeds = 10;
+void run_repeated_reduce_kernel(cuda::stream_ref stream)
+{
+  c2h::device_vector<int> d_out(reuse_repeats, -1);
+  const auto out_ptr = thrust::raw_pointer_cast(d_out.data());
+  const auto config =
+    cuda::make_config(cuda::grid_dims<1>(), cuda::cluster_dims<reuse_cluster_size>(), cuda::block_dims<block_size>());
+  cuda::launch(stream, config, RepeatedReduceKernel{}, out_ptr);
+  stream.sync();
+
+  const c2h::host_vector<int> h_out = d_out;
+  for (int repeat = 0; repeat < reuse_repeats; ++repeat)
+  {
+    const int expected = reuse_cluster_size * block_size * (repeat + 1);
+    REQUIRE(h_out[repeat] == expected);
+  }
+}
+
+inline constexpr int max_size  = 4;
+inline constexpr int num_seeds = 10;
 
 /***********************************************************************************************************************
  * Test cases
@@ -240,3 +279,15 @@ C2H_TEST("reduce/this_cluster Broadcasted", "[reduce][this_cluster]", integral_t
                    c2h::host_vector<value_t>(d_out));
   }
 }
+
+C2H_TEST("reduce/this_cluster can reuse scratch", "[reduce][this_cluster]")
+{
+  const auto device = cuda::devices[0];
+  if (cuda::device_attributes::compute_capability_major(device) < 9)
+  {
+    return;
+  }
+
+  cuda::stream stream{device};
+  run_repeated_reduce_kernel(stream);
+}