From 2dba21564395d9f2c0db8a015cd06b29b284d700 Mon Sep 17 00:00:00 2001 From: Sameer Sahasrabuddhe Date: Thu, 11 Jan 2024 16:19:34 +0530 Subject: [PATCH] SWDEV-295298: introduce warp sync functions The following builtins from the CUDA spec are implemented: - __all_sync, __any_sync, __ballot_sync and __activemask - __match_any_sync and __match_all_sync - __shfl_sync, __shfl_up_sync, __shfl_down_sync, and __shfl_xor_sync The following builtins are NOT implemented, pending support in the compiler: - __reduce_add_sync, __reduce_min_sync, __reduce_max_sync - __reduce_and_sync, __reduce_or_sync, __reduce_xor_sync Change-Id: I07dedbbfe5449f4b5c9b040bed59f5603ccec8c3 [ROCm/clr commit: c5ab5680b400b5b21e9c8c3849ff67e29afcbfc9] --- .../hip/amd_detail/amd_device_functions.h | 34 +-- .../hip/amd_detail/amd_warp_functions.h | 42 +++ .../hip/amd_detail/amd_warp_sync_functions.h | 261 ++++++++++++++++++ projects/clr/hipamd/src/hiprtc/CMakeLists.txt | 3 +- 4 files changed, 306 insertions(+), 34 deletions(-) create mode 100644 projects/clr/hipamd/include/hip/amd_detail/amd_warp_sync_functions.h 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 2a5cf48a03..8d65fd9fde 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 @@ -322,11 +322,6 @@ __device__ static inline unsigned int __usad(unsigned int x, unsigned int y, uns return __ockl_sadd_u32(x, y, z); } -__device__ static inline unsigned int __lane_id() { - return __builtin_amdgcn_mbcnt_hi( - -1, __builtin_amdgcn_mbcnt_lo(-1, 0)); -} - __device__ static inline unsigned int __mbcnt_lo(unsigned int x, unsigned int y) {return __builtin_amdgcn_mbcnt_lo(x,y);}; @@ -339,6 +334,7 @@ HIP specific device functions #if !defined(__HIPCC_RTC__) #include "amd_warp_functions.h" +#include "amd_warp_sync_functions.h" #endif #define MASK1 0x00ff00ff @@ -687,34 +683,6 @@ void __named_sync() { __builtin_amdgcn_s_barrier(); } #endif // __HIP_DEVICE_COMPILE__ -// warp vote function __all __any __ballot -__device__ -inline -int __all(int predicate) { - return __ockl_wfall_i32(predicate); -} - -__device__ -inline -int __any(int predicate) { - return __ockl_wfany_i32(predicate); -} - -// XXX from llvm/include/llvm/IR/InstrTypes.h -#define ICMP_NE 33 - -__device__ -inline -unsigned long long int __ballot(int predicate) { - return __builtin_amdgcn_uicmp(predicate, 0, ICMP_NE); -} - -__device__ -inline -unsigned long long int __ballot64(int predicate) { - return __builtin_amdgcn_uicmp(predicate, 0, ICMP_NE); -} - // hip.amdgcn.bc - lanemask __device__ inline diff --git a/projects/clr/hipamd/include/hip/amd_detail/amd_warp_functions.h b/projects/clr/hipamd/include/hip/amd_detail/amd_warp_functions.h index 559ab20b33..64c2740132 100644 --- a/projects/clr/hipamd/include/hip/amd_detail/amd_warp_functions.h +++ b/projects/clr/hipamd/include/hip/amd_detail/amd_warp_functions.h @@ -75,6 +75,48 @@ __device__ static inline int __hip_move_dpp_N(int src) { static constexpr int warpSize = __AMDGCN_WAVEFRONT_SIZE; +// warp vote function __all __any __ballot +__device__ +inline +int __all(int predicate) { + return __ockl_wfall_i32(predicate); +} + +__device__ +inline +int __any(int predicate) { + return __ockl_wfany_i32(predicate); +} + +// XXX from llvm/include/llvm/IR/InstrTypes.h +#define ICMP_NE 33 + +__device__ +inline +unsigned long long int __ballot(int predicate) { + return __builtin_amdgcn_uicmp(predicate, 0, ICMP_NE); +} + +__device__ +inline +unsigned long long int __ballot64(int predicate) { + return __builtin_amdgcn_uicmp(predicate, 0, ICMP_NE); +} + +// Since threads in a wave do not make independent progress, __activemask() +// always returns the exact active mask, i.e, all active threads in the wave. + +__device__ +inline +unsigned long long __activemask() { + return __ballot(true); +} + +__device__ static inline unsigned int __lane_id() { + return __builtin_amdgcn_mbcnt_hi( + -1, __builtin_amdgcn_mbcnt_lo(-1, 0)); +} + __device__ inline int __shfl(int var, int src_lane, int width = warpSize) { diff --git a/projects/clr/hipamd/include/hip/amd_detail/amd_warp_sync_functions.h b/projects/clr/hipamd/include/hip/amd_detail/amd_warp_sync_functions.h new file mode 100644 index 0000000000..5ce2581a8d --- /dev/null +++ b/projects/clr/hipamd/include/hip/amd_detail/amd_warp_sync_functions.h @@ -0,0 +1,261 @@ +/* +Copyright (c) 2023 Advanced Micro Devices, Inc. All rights reserved. + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in +all copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +THE SOFTWARE. +*/ + +#pragma once + +#if !defined(__HIPCC_RTC__) +#include "amd_warp_functions.h" +#include "hip_assert.h" +#endif + +template +__device__ inline +T __hip_readfirstlane(T val) { + // In theory, behaviour is undefined when reading from a union member other + // than the member that was last assigned to, but it works in practice because + // we rely on the compiler to do the reasonable thing. + union { + unsigned long long l; + T d; + } u; + u.d = val; + // NOTE: The builtin returns int, so we first cast it to unsigned int and only + // then extend it to 64 bits. + unsigned long long lower = (unsigned)__builtin_amdgcn_readfirstlane(u.l); + unsigned long long upper = + (unsigned)__builtin_amdgcn_readfirstlane(u.l >> 32); + u.l = (upper << 32) | lower; + return u.d; +} + +// We use a macro to expand each builtin into a waterfall that implements the +// mask semantics: +// +// 1. The mask argument may be divergent. +// 2. Each active thread must have its own bit set in its own mask value. +// 3. For a given mask value, all threads that are mentioned in the mask must +// execute the same static instance of the builtin with the same mask. +// 4. The union of all mask values supplied at a static instance must be equal +// to the activemask at the program point. +// +// Thus, the mask argument partitions the set of currently active threads in the +// wave into disjoint subsets that cover all active threads. +// +// Implementation notes: +// --------------------- +// +// We implement this as a waterfall loop that executes the builtin for each +// subset separately. The return value is a divergent value across the active +// threads. The value for inactive threads is defined by each builtin +// separately. +// +// As long as every mask value is non-zero, we don't need to check if a lane +// specifies itself in the mask; that is done by the later assertion where all +// chosen lanes must be in the chosen mask. + +#define __hip_check_mask(MASK) \ + do { \ + __hip_assert(MASK && "mask must be non-zero"); \ + bool done = false; \ + while (__any(!done)) { \ + if (!done) { \ + auto chosen_mask = __hip_readfirstlane(MASK); \ + if (MASK == chosen_mask) { \ + __hip_assert(MASK == __ballot(true) && \ + "all threads specified in the mask" \ + " must execute the same operation with the same mask"); \ + done = true; \ + } \ + } \ + } \ + } while(0) + +#define __hip_do_sync(RETVAL, FUNC, MASK, ...) \ + do { \ + __hip_assert(MASK && "mask must be non-zero"); \ + bool done = false; \ + while (__any(!done)) { \ + if (!done) { \ + auto chosen_mask = __hip_readfirstlane(MASK); \ + if (MASK == chosen_mask) { \ + __hip_assert(MASK == __ballot(true) && \ + "all threads specified in the mask" \ + " must execute the same operation with the same mask"); \ + RETVAL = FUNC(__VA_ARGS__); \ + done = true; \ + } \ + } \ + } \ + } while(0) + +// __all_sync, __any_sync, __ballot_sync + +template +__device__ inline +unsigned long long __ballot_sync(MaskT mask, int predicate) { + static_assert( + __hip_internal::is_integral::value && sizeof(MaskT) == 8, + "The mask must be a 64-bit integer. " + "Implicitly promoting a smaller integer is almost always an error."); + __hip_check_mask(mask); + return __ballot(predicate) & mask; +} + +template +__device__ inline +int __all_sync(MaskT mask, int predicate) { + static_assert( + __hip_internal::is_integral::value && sizeof(MaskT) == 8, + "The mask must be a 64-bit integer. " + "Implicitly promoting a smaller integer is almost always an error."); + return __ballot_sync(mask, predicate) == mask; +} + +template +__device__ inline +int __any_sync(MaskT mask, int predicate) { + static_assert( + __hip_internal::is_integral::value && sizeof(MaskT) == 8, + "The mask must be a 64-bit integer. " + "Implicitly promoting a smaller integer is almost always an error."); + return __ballot_sync(mask, predicate) != 0; +} + +// __match_any, __match_all and sync variants + +template +__device__ inline +unsigned long long __match_any(T value) { + static_assert( + (__hip_internal::is_integral::value || __hip_internal::is_floating_point::value) && + (sizeof(T) == 4 || sizeof(T) == 8), + "T can be int, unsigned int, long, unsigned long, long long, unsigned " + "long long, float or double."); + bool done = false; + unsigned long long retval = 0; + + while (__any(!done)) { + if (!done) { + T chosen = __hip_readfirstlane(value); + if (chosen == value) { + retval = __activemask(); + done = true; + } + } + } + + return retval; +} + +template +__device__ inline +unsigned long long __match_any_sync(MaskT mask, T value) { + static_assert( + __hip_internal::is_integral::value && sizeof(MaskT) == 8, + "The mask must be a 64-bit integer. " + "Implicitly promoting a smaller integer is almost always an error."); + __hip_check_mask(mask); + return __match_any(value) & mask; +} + +template +__device__ inline +unsigned long long __match_all(T value, int* pred) { + static_assert( + (__hip_internal::is_integral::value || __hip_internal::is_floating_point::value) && + (sizeof(T) == 4 || sizeof(T) == 8), + "T can be int, unsigned int, long, unsigned long, long long, unsigned " + "long long, float or double."); + T first = __hip_readfirstlane(value); + if (__all(first == value)) { + *pred = true; + return __activemask(); + } else { + *pred = false; + return 0; + } +} + +template +__device__ inline +unsigned long long __match_all_sync(MaskT mask, T value, int* pred) { + static_assert( + __hip_internal::is_integral::value && sizeof(MaskT) == 8, + "The mask must be a 64-bit integer. " + "Implicitly promoting a smaller integer is almost always an error."); + MaskT retval = 0; + __hip_do_sync(retval, __match_all, mask, value, pred); + return retval; +} + +// various variants of shfl + +template +__device__ inline +T __shfl_sync(MaskT mask, T var, int srcLane, + int width = __AMDGCN_WAVEFRONT_SIZE) { + static_assert( + __hip_internal::is_integral::value && sizeof(MaskT) == 8, + "The mask must be a 64-bit integer. " + "Implicitly promoting a smaller integer is almost always an error."); + __hip_check_mask(mask); + return __shfl(var, srcLane, width); +} + +template +__device__ inline +T __shfl_up_sync(MaskT mask, T var, unsigned int delta, + int width = __AMDGCN_WAVEFRONT_SIZE) { + static_assert( + __hip_internal::is_integral::value && sizeof(MaskT) == 8, + "The mask must be a 64-bit integer. " + "Implicitly promoting a smaller integer is almost always an error."); + __hip_check_mask(mask); + return __shfl_up(var, delta, width); +} + +template +__device__ inline +T __shfl_down_sync(MaskT mask, T var, unsigned int delta, + int width = __AMDGCN_WAVEFRONT_SIZE) { + static_assert( + __hip_internal::is_integral::value && sizeof(MaskT) == 8, + "The mask must be a 64-bit integer. " + "Implicitly promoting a smaller integer is almost always an error."); + __hip_check_mask(mask); + return __shfl_down(var, delta, width); +} + +template +__device__ inline +T __shfl_xor_sync(MaskT mask, T var, int laneMask, + int width = __AMDGCN_WAVEFRONT_SIZE) { + static_assert( + __hip_internal::is_integral::value && sizeof(MaskT) == 8, + "The mask must be a 64-bit integer. " + "Implicitly promoting a smaller integer is almost always an error."); + __hip_check_mask(mask); + return __shfl_xor(var, laneMask, width); +} + +#undef __hip_do_sync +#undef __hip_check_mask diff --git a/projects/clr/hipamd/src/hiprtc/CMakeLists.txt b/projects/clr/hipamd/src/hiprtc/CMakeLists.txt index e0b0e9417f..b1745a02f5 100644 --- a/projects/clr/hipamd/src/hiprtc/CMakeLists.txt +++ b/projects/clr/hipamd/src/hiprtc/CMakeLists.txt @@ -146,8 +146,9 @@ ${PROJECT_SOURCE_DIR}/include/hip/amd_detail/amd_surface_functions.h ${PROJECT_SOURCE_DIR}/include/hip/amd_detail/amd_hip_complex.h ${PROJECT_SOURCE_DIR}/include/hip/amd_detail/amd_hip_math_constants.h ${PROJECT_SOURCE_DIR}/include/hip/amd_detail/math_fwd.h -${PROJECT_SOURCE_DIR}/include/hip/amd_detail/amd_device_functions.h ${PROJECT_SOURCE_DIR}/include/hip/amd_detail/amd_warp_functions.h +${PROJECT_SOURCE_DIR}/include/hip/amd_detail/amd_device_functions.h +${PROJECT_SOURCE_DIR}/include/hip/amd_detail/amd_warp_sync_functions.h ${PROJECT_SOURCE_DIR}/include/hip/amd_detail/hip_cooperative_groups_helper.h ${PROJECT_SOURCE_DIR}/include/hip/amd_detail/amd_hip_cooperative_groups.h ${PROJECT_SOURCE_DIR}/include/hip/amd_detail/amd_hip_unsafe_atomics.h