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 ca9dae4437..2661893a59 100644 --- a/hipamd/include/hip/amd_detail/amd_hip_cooperative_groups.h +++ b/hipamd/include/hip/amd_detail/amd_hip_cooperative_groups.h @@ -91,6 +91,8 @@ class thread_group { struct _tiled_info { bool is_tiled; unsigned int size; + unsigned int meta_group_rank; + unsigned int meta_group_size; }; struct _coalesced_info { @@ -225,7 +227,6 @@ class thread_block : public thread_group { unsigned int tile_size); friend __CG_QUALIFIER__ thread_group tiled_partition(const thread_block& parent, unsigned int tile_size); - protected: // Construct a workgroup thread group (through the API this_thread_block()) explicit __CG_QUALIFIER__ thread_block(uint32_t size) @@ -235,23 +236,26 @@ class thread_block : public thread_group { const bool pow2 = ((tile_size & (tile_size - 1)) == 0); // Invalid tile size, assert if (!tile_size || (tile_size > __AMDGCN_WAVEFRONT_SIZE) || !pow2) { - __hip_assert(false && "invalid tile size") + __hip_assert(false && "invalid tile size"); } thread_group tiledGroup = thread_group(internal::cg_tiled_group, tile_size); tiledGroup.coalesced_info.tiled_info.size = tile_size; tiledGroup.coalesced_info.tiled_info.is_tiled = true; + tiledGroup.coalesced_info.tiled_info.meta_group_rank = thread_rank() / tile_size; + tiledGroup.coalesced_info.tiled_info.meta_group_size = (size() + tile_size - 1) / tile_size; return tiledGroup; } public: // 3-dimensional block index within the grid - __CG_QUALIFIER__ dim3 group_index() { return internal::workgroup::group_index(); } + __CG_STATIC_QUALIFIER__ dim3 group_index() { return internal::workgroup::group_index(); } // 3-dimensional thread index within the block - __CG_QUALIFIER__ dim3 thread_index() { return internal::workgroup::thread_index(); } - __CG_QUALIFIER__ uint32_t thread_rank() const { return internal::workgroup::thread_rank(); } - __CG_QUALIFIER__ bool is_valid() const { return internal::workgroup::is_valid(); } - __CG_QUALIFIER__ void sync() const { internal::workgroup::sync(); } + __CG_STATIC_QUALIFIER__ dim3 thread_index() { return internal::workgroup::thread_index(); } + __CG_STATIC_QUALIFIER__ uint32_t thread_rank() { return internal::workgroup::thread_rank(); } + __CG_STATIC_QUALIFIER__ uint32_t size() { return internal::workgroup::size(); } + __CG_STATIC_QUALIFIER__ bool is_valid() { return internal::workgroup::is_valid(); } + __CG_STATIC_QUALIFIER__ void sync() { internal::workgroup::sync(); } __CG_QUALIFIER__ dim3 group_dim() { return internal::workgroup::block_dim(); } }; @@ -287,7 +291,7 @@ class tiled_group : public thread_group { const bool pow2 = ((tile_size & (tile_size - 1)) == 0); if (!tile_size || (tile_size > __AMDGCN_WAVEFRONT_SIZE) || !pow2) { - __hip_assert(false && "invalid tile size") + __hip_assert(false && "invalid tile size"); } if (size() <= tile_size) { @@ -348,6 +352,8 @@ class coalesced_group : public thread_group { member_mask <<= (__lane_id() & ~(tile_size - 1)); coalesced_group coalesced_tile = coalesced_group(member_mask); coalesced_tile.coalesced_info.tiled_info.is_tiled = true; + coalesced_tile.coalesced_info.tiled_info.meta_group_rank = thread_rank() / tile_size; + coalesced_tile.coalesced_info.tiled_info.meta_group_size = size() / tile_size; return coalesced_tile; } // Here the parent coalesced_group is not partitioned. @@ -369,6 +375,9 @@ class coalesced_group : public thread_group { } } coalesced_group coalesced_tile = coalesced_group(member_mask); + coalesced_tile.coalesced_info.tiled_info.meta_group_rank = thread_rank() / tile_size; + coalesced_tile.coalesced_info.tiled_info.meta_group_size = + (size() + tile_size - 1) / tile_size; return coalesced_tile; } return coalesced_group(0); @@ -396,6 +405,14 @@ class coalesced_group : public thread_group { internal::coalesced_group::sync(); } + __CG_QUALIFIER__ unsigned int meta_group_rank() const { + return coalesced_info.tiled_info.meta_group_rank; + } + + __CG_QUALIFIER__ unsigned int meta_group_size() const { + return coalesced_info.tiled_info.meta_group_size; + } + template __CG_QUALIFIER__ T shfl(T var, int srcRank) const { static_assert(is_valid_type::value, "Neither an integer or float type."); @@ -498,7 +515,7 @@ __CG_QUALIFIER__ uint32_t thread_group::thread_rank() const { return (static_cast(this)->thread_rank()); } default: { - __hip_assert(false && "invalid cooperative group type") + __hip_assert(false && "invalid cooperative group type"); return -1; } } @@ -526,7 +543,7 @@ __CG_QUALIFIER__ bool thread_group::is_valid() const { return (static_cast(this)->is_valid()); } default: { - __hip_assert(false && "invalid cooperative group type") + __hip_assert(false && "invalid cooperative group type"); return false; } } @@ -559,7 +576,7 @@ __CG_QUALIFIER__ void thread_group::sync() const { break; } default: { - __hip_assert(false && "invalid cooperative group type") + __hip_assert(false && "invalid cooperative group type"); } } } @@ -647,6 +664,22 @@ template class thread_block_tile_base : public tile_base +class parent_group_info { +public: + // Returns the linear rank of the group within the set of tiles partitioned + // from a parent group (bounded by meta_group_size) + __CG_STATIC_QUALIFIER__ unsigned int meta_group_rank() { + return ParentCGTy::thread_rank() / tileSize; + } + + // Returns the number of groups created when the parent group was partitioned. + __CG_STATIC_QUALIFIER__ unsigned int meta_group_size() { + return (ParentCGTy::size() + tileSize - 1) / tileSize; + } +}; /** \brief Group type - thread_block_tile * @@ -654,25 +687,49 @@ template class thread_block_tile_base : public tile_base -class thread_block_tile_type : public thread_block_tile_base, public tiled_group { +template +class thread_block_tile_type : public thread_block_tile_base, + public tiled_group, + public parent_group_info { _CG_STATIC_CONST_DECL_ unsigned int numThreads = tileSize; + protected: + __CG_QUALIFIER__ thread_block_tile_type() : tiled_group(numThreads) { + coalesced_info.tiled_info.size = numThreads; + coalesced_info.tiled_info.is_tiled = true; + } +}; - friend class thread_block_tile_type; +// Partial template specialization +template +class thread_block_tile_type : public thread_block_tile_base, + public tiled_group + { + _CG_STATIC_CONST_DECL_ unsigned int numThreads = tileSize; typedef thread_block_tile_base tbtBase; protected: - __CG_QUALIFIER__ thread_block_tile_type() : tiled_group(numThreads) { + + __CG_QUALIFIER__ thread_block_tile_type(unsigned int meta_group_rank, unsigned int meta_group_size) + : tiled_group(numThreads) { coalesced_info.tiled_info.size = numThreads; coalesced_info.tiled_info.is_tiled = true; + coalesced_info.tiled_info.meta_group_rank = meta_group_rank; + coalesced_info.tiled_info.meta_group_size = meta_group_size; } public: using tbtBase::size; using tbtBase::sync; using tbtBase::thread_rank; + + __CG_QUALIFIER__ unsigned int meta_group_rank() const { + return coalesced_info.tiled_info.meta_group_rank; + } + + __CG_QUALIFIER__ unsigned int meta_group_size() const { + return coalesced_info.tiled_info.meta_group_size; + } // end of operative group /** * @} @@ -726,7 +783,7 @@ class thread_block_tile_internal : public thread_block_tile_type __CG_QUALIFIER__ thread_block_tile_internal( const thread_block_tile_internal& g) - : thread_block_tile_type() {} + : thread_block_tile_type(g.meta_group_rank(), g.meta_group_size()) {} __CG_QUALIFIER__ thread_block_tile_internal(const thread_block& g) : thread_block_tile_type() {} @@ -760,7 +817,7 @@ class thread_block_tile : public impl::thread_block_tile_internal class thread_block_tile; namespace impl { -template struct tiled_partition_internal; +template struct tiled_partition_internal; template struct tiled_partition_internal : public thread_block_tile {