From 40e2603bd08b187da35cda99b219251d64c767fb Mon Sep 17 00:00:00 2001 From: Sourabh Betigeri Date: Thu, 18 Nov 2021 17:54:46 -0800 Subject: [PATCH] SWDEV-294375 - Adding coalesced_groups to ROCm CG Change-Id: I4098e1970ae2c388171fa73deb5bc023ed1b7a68 [ROCm/clr commit: 523e5a007c38cdb91088b9b97a33fc149f7eac90] --- .../hip/amd_detail/amd_device_functions.h | 72 ++++++ .../amd_detail/amd_hip_cooperative_groups.h | 207 ++++++++++++++++-- .../hip_cooperative_groups_helper.h | 31 ++- 3 files changed, 294 insertions(+), 16 deletions(-) diff --git a/projects/clr/hipamd/include/hip/amd_detail/amd_device_functions.h b/projects/clr/hipamd/include/hip/amd_detail/amd_device_functions.h index f9b9570dfc..e0b5ee0774 100644 --- a/projects/clr/hipamd/include/hip/amd_detail/amd_device_functions.h +++ b/projects/clr/hipamd/include/hip/amd_detail/amd_device_functions.h @@ -78,6 +78,78 @@ __device__ static inline unsigned int __ffsll(long long int input) { return ( input == 0 ? -1 : __builtin_ctzll(input) ) + 1; } +// Given a 32/64-bit value exec mask and an integer value base (between 0 and WAVEFRONT_SIZE), +// find the n-th (given by offset) set bit in the exec mask from the base bit, and return the bit position. +// If not found, return -1. +__device__ int32_t __fns64(uint64_t mask, uint32_t base, int32_t offset) { + uint64_t temp_mask = mask; + int32_t temp_offset = offset; + + if (offset == 0) { + temp_mask &= (1 << base); + temp_offset = 1; + } + else if (offset < 0) { + temp_mask = __builtin_bitreverse64(mask); + base = 63 - base; + temp_offset = -offset; + } + + temp_mask = temp_mask & ((~0ULL) << base); + if (__builtin_popcountll(temp_mask) < temp_offset) + return -1; + int32_t total = 0; + for (int i = 0x20; i > 0; i >>= 1) { + uint64_t temp_mask_lo = temp_mask & ((1ULL << i) - 1); + uint32_t pcnt = __builtin_popcountll(temp_mask_lo); + if (pcnt < temp_offset) { + temp_mask = temp_mask >> i; + temp_offset -= pcnt; + total += i; + } + else { + temp_mask = temp_mask_lo; + } + } + if (offset < 0) + return 63 - total; + else + return total; +} + +__device__ int32_t __fns32(uint64_t mask, uint32_t base, int32_t offset) { + uint64_t temp_mask = mask; + int32_t temp_offset = offset; + if (offset == 0) { + temp_mask &= (1 << base); + temp_offset = 1; + } + else if (offset < 0) { + temp_mask = __builtin_bitreverse64(mask); + base = 63 - base; + temp_offset = -offset; + } + temp_mask = temp_mask & ((~0ULL) << base); + if (__builtin_popcountll(temp_mask) < temp_offset) + return -1; + int32_t total = 0; + for (int i = 0x20; i > 0; i >>= 1) { + uint64_t temp_mask_lo = temp_mask & ((1ULL << i) - 1); + uint32_t pcnt = __builtin_popcountll(temp_mask_lo); + if (pcnt < temp_offset) { + temp_mask = temp_mask >> i; + temp_offset -= pcnt; + total += i; + } + else { + temp_mask = temp_mask_lo; + } + } + if (offset < 0) + return 63 - total; + else + return total; +} __device__ static inline unsigned int __brev(unsigned int input) { return __builtin_bitreverse32(input); } diff --git a/projects/clr/hipamd/include/hip/amd_detail/amd_hip_cooperative_groups.h b/projects/clr/hipamd/include/hip/amd_detail/amd_hip_cooperative_groups.h index 65cd8508b6..03aa4d3297 100644 --- a/projects/clr/hipamd/include/hip/amd_detail/amd_hip_cooperative_groups.h +++ b/projects/clr/hipamd/include/hip/amd_detail/amd_hip_cooperative_groups.h @@ -54,7 +54,7 @@ class thread_group { // only when the group is supposed to contain only the calling the thread // (throurh the API - `this_thread()`), and in all other cases, this thread // group object is a sub-object of some other derived thread group object - __CG_QUALIFIER__ thread_group(internal::group_type type, uint32_t size, + __CG_QUALIFIER__ thread_group(internal::group_type type, uint32_t size = (uint64_t)0, uint64_t mask = (uint64_t)0) { _type = type; _size = size; @@ -64,7 +64,13 @@ class thread_group { struct _tiled_info { bool is_tiled; unsigned int size; - } tiled_info; + }; + + struct _coalesced_info { + lane_mask member_mask; + unsigned int size; + struct _tiled_info tiled_info; + } coalesced_info; friend __CG_QUALIFIER__ thread_group tiled_partition(const thread_group& parent, unsigned int tile_size); @@ -182,8 +188,8 @@ class thread_block : public thread_group { } thread_group tiledGroup = thread_group(internal::cg_tiled_group, tile_size); - tiledGroup.tiled_info.size = tile_size; - tiledGroup.tiled_info.is_tiled = true; + tiledGroup.coalesced_info.tiled_info.size = tile_size; + tiledGroup.coalesced_info.tiled_info.is_tiled = true; return tiledGroup; } @@ -229,26 +235,26 @@ class tiled_group : public thread_group { } if (size() <= tile_size) { - return (*this); + return *this; } tiled_group tiledGroup = tiled_group(tile_size); - tiledGroup.tiled_info.is_tiled = true; + tiledGroup.coalesced_info.tiled_info.is_tiled = true; return tiledGroup; } protected: explicit __CG_QUALIFIER__ tiled_group(unsigned int tileSize) : thread_group(internal::cg_tiled_group, tileSize) { - tiled_info.size = tileSize; - tiled_info.is_tiled = true; + coalesced_info.tiled_info.size = tileSize; + coalesced_info.tiled_info.is_tiled = true; } public: - __CG_QUALIFIER__ unsigned int size() const { return (tiled_info.size); } + __CG_QUALIFIER__ unsigned int size() const { return (coalesced_info.tiled_info.size); } __CG_QUALIFIER__ unsigned int thread_rank() const { - return (internal::workgroup::thread_rank() & (tiled_info.size - 1)); + return (internal::workgroup::thread_rank() & (coalesced_info.tiled_info.size - 1)); } __CG_QUALIFIER__ void sync() const { @@ -256,6 +262,160 @@ class tiled_group : public thread_group { } }; +/** \brief The coalesced_group cooperative group type + * + * \details Represents a active thread group in a wavefront. + * This group type also supports sub-wave level intrinsics. + */ + +class coalesced_group : public thread_group { + private: + 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); + + __CG_QUALIFIER__ coalesced_group new_tiled_group(unsigned int tile_size) const { + const bool pow2 = ((tile_size & (tile_size - 1)) == 0); + + if (!tile_size || (tile_size > size()) || !pow2) { + return coalesced_group(0); + } + + // If a tiled group is passed to be partitioned further into a coalesced_group. + // prepare a mask for further partitioning it so that it stays coalesced. + if (coalesced_info.tiled_info.is_tiled) { + unsigned int base_offset = (thread_rank() & (~(tile_size - 1))); + unsigned int masklength = min((unsigned int)size() - base_offset, tile_size); + lane_mask member_mask = (lane_mask)(-1) >> (WAVEFRONT_SIZE - masklength); + + member_mask <<= (__lane_id() & ~(tile_size - 1)); + coalesced_group coalesced_tile = coalesced_group(member_mask); + coalesced_tile.coalesced_info.tiled_info.is_tiled = true; + return coalesced_tile; + } + // Here the parent coalesced_group is not partitioned. + else { + lane_mask member_mask = 0; + unsigned int tile_rank = 0; + int lanes_to_skip = ((thread_rank()) / tile_size) * tile_size; + + for (unsigned int i = 0; i < WAVEFRONT_SIZE; i++) { + lane_mask active = coalesced_info.member_mask & (1 << i); + // Make sure the lane is active + if (active) { + if (lanes_to_skip <= 0 && tile_rank < tile_size) { + // Prepare a member_mask that is appropriate for a tile + member_mask |= active; + tile_rank++; + } + lanes_to_skip--; + } + } + coalesced_group coalesced_tile = coalesced_group(member_mask); + return coalesced_tile; + } + return coalesced_group(0); + } + + protected: + // Constructor + explicit __CG_QUALIFIER__ coalesced_group(lane_mask member_mask) + : thread_group(internal::cg_coalesced_group) { + coalesced_info.member_mask = member_mask; // Which threads are active + coalesced_info.size = __popcll(coalesced_info.member_mask); // How many threads are active + coalesced_info.tiled_info.is_tiled = false; // Not a partitioned group + } + + public: + __CG_QUALIFIER__ unsigned int size() const { + return coalesced_info.size; + } + + __CG_QUALIFIER__ unsigned int thread_rank() const { + return internal::coalesced_group::masked_bit_count(coalesced_info.member_mask); + } + + __CG_QUALIFIER__ void sync() const { + internal::coalesced_group::sync(); + } + + template + __CG_QUALIFIER__ T shfl(T var, int srcRank) const { + static_assert(is_valid_type::value, "Neither an integer or float type."); + + srcRank = srcRank % size(); + + int lane = (size() == WAVEFRONT_SIZE) ? srcRank + : (WAVEFRONT_SIZE == 64) ? __fns64(coalesced_info.member_mask, 0, (srcRank + 1)) + : __fns32(coalesced_info.member_mask, 0, (srcRank + 1)); + + return __shfl(var, lane, WAVEFRONT_SIZE); + } + + template + __CG_QUALIFIER__ T shfl_down(T var, unsigned int lane_delta) const { + static_assert(is_valid_type::value, "Neither an integer or float type."); + + // Note: The cuda implementation appears to use the remainder of lane_delta + // and WARP_SIZE as the shift value rather than lane_delta itself. + // This is not described in the documentation and is not done here. + + if (size() == WAVEFRONT_SIZE) { + return __shfl_down(var, lane_delta, WAVEFRONT_SIZE); + } + + int lane; + if (WAVEFRONT_SIZE == 64) { + lane = __fns64(coalesced_info.member_mask, __lane_id(), lane_delta + 1); + } + else { + lane = __fns32(coalesced_info.member_mask, __lane_id(), lane_delta + 1); + } + + if (lane == -1) { + lane = __lane_id(); + } + + return __shfl(var, lane, WAVEFRONT_SIZE); + } + + template + __CG_QUALIFIER__ T shfl_up(T var, unsigned int lane_delta) const { + static_assert(is_valid_type::value, "Neither an integer or float type."); + + // Note: The cuda implementation appears to use the remainder of lane_delta + // and WARP_SIZE as the shift value rather than lane_delta itself. + // This is not described in the documentation and is not done here. + + if (size() == WAVEFRONT_SIZE) { + return __shfl_up(var, lane_delta, WAVEFRONT_SIZE); + } + + int lane; + if (WAVEFRONT_SIZE == 64) { + lane = __fns64(coalesced_info.member_mask, __lane_id(), -(lane_delta + 1)); + } + else if (WAVEFRONT_SIZE == 32) { + lane = __fns32(coalesced_info.member_mask, __lane_id(), -(lane_delta + 1)); + } + + if (lane == -1) { + lane = __lane_id(); + } + + return __shfl(var, lane, WAVEFRONT_SIZE); + } +}; + +/** \brief User exposed API to create coalesced groups. + * + * \details A collective operation that groups all active lanes into a new thread group. + */ + +__CG_QUALIFIER__ coalesced_group coalesced_threads() { + return cooperative_groups::coalesced_group(__builtin_amdgcn_read_exec()); +} + /** * Implemenation of all publicly exposed base class APIs */ @@ -273,6 +433,9 @@ __CG_QUALIFIER__ uint32_t thread_group::thread_rank() const { case internal::cg_tiled_group: { return (static_cast(this)->thread_rank()); } + case internal::cg_coalesced_group: { + return (static_cast(this)->thread_rank()); + } default: { assert(false && "invalid cooperative group type"); return -1; @@ -294,6 +457,9 @@ __CG_QUALIFIER__ bool thread_group::is_valid() const { case internal::cg_tiled_group: { return (static_cast(this)->is_valid()); } + case internal::cg_coalesced_group: { + return (static_cast(this)->is_valid()); + } default: { assert(false && "invalid cooperative group type"); return false; @@ -319,6 +485,10 @@ __CG_QUALIFIER__ void thread_group::sync() const { static_cast(this)->sync(); break; } + case internal::cg_coalesced_group: { + static_cast(this)->sync(); + break; + } default: { assert(false && "invalid cooperative group type"); } @@ -399,8 +569,8 @@ class thread_block_tile_type : public thread_block_tile_base, public t protected: __CG_QUALIFIER__ thread_block_tile_type() : tiled_group(numThreads) { - tiled_info.size = numThreads; - tiled_info.is_tiled = true; + coalesced_info.tiled_info.size = numThreads; + coalesced_info.tiled_info.is_tiled = true; } public: @@ -420,7 +590,12 @@ __CG_QUALIFIER__ thread_group tiled_partition(const thread_group& parent, unsign if (parent.cg_type() == internal::cg_tiled_group) { const tiled_group* cg = static_cast(&parent); return cg->new_tiled_group(tile_size); - } else { + } + else if(parent.cg_type() == internal::cg_coalesced_group) { + const coalesced_group* cg = static_cast(&parent); + return cg->new_tiled_group(tile_size); + } + else { const thread_block* tb = static_cast(&parent); return tb->new_tiled_group(tile_size); } @@ -431,11 +606,15 @@ __CG_QUALIFIER__ thread_group tiled_partition(const thread_block& parent, unsign return (parent.new_tiled_group(tile_size)); } -// Coalesced group type overload __CG_QUALIFIER__ tiled_group tiled_partition(const tiled_group& parent, unsigned int tile_size) { return (parent.new_tiled_group(tile_size)); } +// If a coalesced group is passed to be partitioned, it should remain coalesced +__CG_QUALIFIER__ coalesced_group tiled_partition(const coalesced_group& parent, unsigned int tile_size) { + return (parent.new_tiled_group(tile_size)); +} + template class thread_block_tile; namespace impl { diff --git a/projects/clr/hipamd/include/hip/amd_detail/hip_cooperative_groups_helper.h b/projects/clr/hipamd/include/hip/amd_detail/hip_cooperative_groups_helper.h index e10ae3b9d3..eaee1a715a 100644 --- a/projects/clr/hipamd/include/hip/amd_detail/hip_cooperative_groups_helper.h +++ b/projects/clr/hipamd/include/hip/amd_detail/hip_cooperative_groups_helper.h @@ -33,7 +33,7 @@ THE SOFTWARE. #if __cplusplus #include - +#include #if !defined(__align__) #define __align__(x) __attribute__((aligned(x))) #endif @@ -53,8 +53,10 @@ THE SOFTWARE. #if !defined(WAVEFRONT_SIZE) #if __gfx1010__ || __gfx1011__ || __gfx1012__ || __gfx1030__ || __gfx1031__ #define WAVEFRONT_SIZE 32 +using lane_mask = unsigned int; #else #define WAVEFRONT_SIZE 64 +using lane_mask = unsigned long long int; #endif namespace cooperative_groups { @@ -83,7 +85,8 @@ typedef enum { cg_multi_grid, cg_grid, cg_workgroup, - cg_tiled_group + cg_tiled_group, + cg_coalesced_group } group_type; /** @@ -177,6 +180,30 @@ __CG_STATIC_QUALIFIER__ void sync() { __builtin_amdgcn_fence(__ATOMIC_ACQ_REL, " } // namespace tiled_group +namespace coalesced_group { + +// enforce ordering for memory intructions +__CG_STATIC_QUALIFIER__ void sync() { __builtin_amdgcn_fence(__ATOMIC_ACQ_REL, "agent"); } + +// Masked bit count +// +// For each thread, this function returns the number of active threads which +// have i-th bit of x set and come before the current thread. +__device__ unsigned int masked_bit_count(lane_mask x, unsigned int add = 0) { + int counter=0; + #if WAVEFRONT_SIZE == 32 + counter = __builtin_amdgcn_mbcnt_lo(x, add); + #else + counter = __builtin_amdgcn_mbcnt_lo(static_cast(x), add); + counter = __builtin_amdgcn_mbcnt_hi(static_cast(x >> 32), counter); + #endif + + return counter; +} + +} // namespace coalesced_group + + } // namespace internal } // namespace cooperative_groups