diff --git a/hipamd/include/hip/amd_detail/amd_hip_cooperative_groups.h b/hipamd/include/hip/amd_detail/amd_hip_cooperative_groups.h index c01039a7e1..6c038d33ec 100644 --- a/hipamd/include/hip/amd_detail/amd_hip_cooperative_groups.h +++ b/hipamd/include/hip/amd_detail/amd_hip_cooperative_groups.h @@ -79,6 +79,7 @@ class thread_group { struct _tiled_info tiled_info; } coalesced_info; + friend __CG_QUALIFIER__ thread_group this_thread(); friend __CG_QUALIFIER__ thread_group tiled_partition(const thread_group& parent, unsigned int tile_size); friend class thread_block; @@ -175,6 +176,7 @@ class grid_group : public thread_group { __CG_QUALIFIER__ uint32_t thread_rank() const { return internal::grid::thread_rank(); } __CG_QUALIFIER__ bool is_valid() const { return internal::grid::is_valid(); } __CG_QUALIFIER__ void sync() const { internal::grid::sync(); } + __CG_QUALIFIER__ dim3 group_dim() const { return internal::workgroup::block_dim(); } }; /** @brief User exposed API interface to construct grid cooperative group type @@ -306,6 +308,8 @@ class tiled_group : public thread_group { } }; +template class thread_block_tile; + /** \brief The coalesced_group cooperative group type * * \details Represents a active thread group in a wavefront. @@ -318,6 +322,10 @@ class coalesced_group : public thread_group { friend __CG_QUALIFIER__ coalesced_group coalesced_threads(); friend __CG_QUALIFIER__ thread_group tiled_partition(const thread_group& parent, unsigned int tile_size); friend __CG_QUALIFIER__ coalesced_group tiled_partition(const coalesced_group& parent, unsigned int tile_size); + friend __CG_QUALIFIER__ coalesced_group binary_partition(const coalesced_group& cgrp, bool pred); + template + friend __CG_QUALIFIER__ coalesced_group + binary_partition(const thread_block_tile& tgrp, bool pred); __CG_QUALIFIER__ coalesced_group new_tiled_group(unsigned int tile_size) const { const bool pow2 = ((tile_size & (tile_size - 1)) == 0); @@ -465,6 +473,34 @@ class coalesced_group : public thread_group { return __shfl(var, lane, __AMDGCN_WAVEFRONT_SIZE); } +#ifdef HIP_ENABLE_WARP_SYNC_BUILTINS + __CG_QUALIFIER__ unsigned long long ballot(int pred) const { + return internal::helper::adjust_mask( + coalesced_info.member_mask, + __ballot_sync(coalesced_info.member_mask, pred)); + } + + __CG_QUALIFIER__ int any(int pred) const { + return __any_sync(static_cast(coalesced_info.member_mask), pred); + } + + __CG_QUALIFIER__ int all(int pred) const { + return __all_sync(static_cast(coalesced_info.member_mask), pred); + } + + template __CG_QUALIFIER__ unsigned long long match_any(T value) const { + return internal::helper::adjust_mask( + coalesced_info.member_mask, + __match_any_sync(static_cast(coalesced_info.member_mask), value)); + } + + template __CG_QUALIFIER__ unsigned long long match_all(T value, int& pred) const { + return internal::helper::adjust_mask( + coalesced_info.member_mask, + __match_all_sync(static_cast(coalesced_info.member_mask), value, + &pred)); + } +#endif }; /** \brief User exposed API to create coalesced groups. @@ -625,6 +661,17 @@ template class thread_block_tile_base : public tile_base::numThreads; + template + friend __CG_QUALIFIER__ coalesced_group + binary_partition(const thread_block_tile& tgrp, bool pred); + +#ifdef HIP_ENABLE_WARP_SYNC_BUILTINS + __CG_QUALIFIER__ unsigned long long build_mask() const { + unsigned long long mask = ~0ull >> (64 - numThreads); + return mask << ((internal::workgroup::thread_rank() / numThreads) * numThreads); + } +#endif + public: __CG_STATIC_QUALIFIER__ void sync() { internal::tiled_group::sync(); @@ -649,7 +696,29 @@ template class thread_block_tile_base : public tile_base::value, "Neither an integer or float type."); return (__shfl_xor(var, laneMask, numThreads)); } + +#ifdef HIP_ENABLE_WARP_SYNC_BUILTINS + __CG_QUALIFIER__ unsigned long long ballot(int pred) const { + const auto mask = build_mask(); + return internal::helper::adjust_mask(mask, __ballot_sync(mask, pred)); + } + + __CG_QUALIFIER__ int any(int pred) const { return __any_sync(build_mask(), pred); } + + __CG_QUALIFIER__ int all(int pred) const { return __all_sync(build_mask(), pred); } + + template __CG_QUALIFIER__ unsigned long long match_any(T value) const { + const auto mask = build_mask(); + return internal::helper::adjust_mask(mask, __match_any_sync(mask, value)); + } + + template __CG_QUALIFIER__ unsigned long long match_all(T value, int& pred) const { + const auto mask = build_mask(); + return internal::helper::adjust_mask(mask, __match_all_sync(mask, value, &pred)); + } +#endif }; + /** \brief User exposed API that captures the state of the parent group pre-partition */ template @@ -727,6 +796,10 @@ class thread_block_tile_type : public thread_block_tile_base class thread_block_tile; - namespace impl { template class thread_block_tile_internal; @@ -829,6 +900,34 @@ __CG_QUALIFIER__ thread_block_tile tiled_partition(const Paren "Tiled partition with size > wavefront size. Currently not supported "); return impl::tiled_partition_internal(g); } + +#ifdef HIP_ENABLE_WARP_SYNC_BUILTINS +/** \brief Binary partition + * + * \details This splits the input thread group into two partitions determined by predicate + */ +__CG_QUALIFIER__ coalesced_group binary_partition(const coalesced_group& cgrp, bool pred) { + auto mask = __ballot_sync(cgrp.coalesced_info.member_mask, pred); + + if (pred) { + return coalesced_group(mask); + } else { + return coalesced_group(cgrp.coalesced_info.member_mask ^ mask); + } +} + +template +__CG_QUALIFIER__ coalesced_group binary_partition(const thread_block_tile& tgrp, + bool pred) { + auto mask = __ballot_sync(tgrp.build_mask(), pred); + + if (pred) { + return coalesced_group(mask); + } else { + return coalesced_group(tgrp.build_mask() ^ mask); + } +} +#endif } // namespace cooperative_groups #endif // __cplusplus diff --git a/hipamd/include/hip/amd_detail/hip_cooperative_groups_helper.h b/hipamd/include/hip/amd_detail/hip_cooperative_groups_helper.h index f62246a031..95379bcaef 100644 --- a/hipamd/include/hip/amd_detail/hip_cooperative_groups_helper.h +++ b/hipamd/include/hip/amd_detail/hip_cooperative_groups_helper.h @@ -102,6 +102,33 @@ typedef enum { * on Windows. * */ +namespace helper { +/** + * @brief Create output mask from input_mask at places where base_mask is set + * + * Example: base_mask = 0101'0101, input_mask = 1111'0000 + * Output mask: 1100 + * Explaination: + * | | | | | | | | + * base: 0|1|0|1|'0|1|0|1| // Which bits are set + * input: 1|1|1|1|'0|0|0|0| // Which values are picked + * | | | | | | | | + * output: 1 1 0 0 + */ +__CG_STATIC_QUALIFIER__ unsigned long long adjust_mask( + unsigned long long base_mask, unsigned long long input_mask) { + unsigned long long out = 0; + for (unsigned int i = 0, index = 0; i < __AMDGCN_WAVEFRONT_SIZE; i++) { + auto lane_active = base_mask & (1ull << i); + if (lane_active) { + auto result = input_mask & (1ull << i); + out |= ((result ? 1ull : 0ull) << index); + index++; + } + } + return out; +} +} // namespace helper /** * * @brief Functionalities related to multi-grid cooperative group type