From 340d6bb69f1846086b888eea28fb8919972ae66b Mon Sep 17 00:00:00 2001 From: Gerardo Hernandez Date: Mon, 25 Nov 2024 14:32:40 +0000 Subject: [PATCH] SWDEV-420237 - Add __reduce_add_sync() Change-Id: Ic8e4fab6b7aeb879d40b2c1419b30d1355a2bbdc --- .../hip/amd_detail/amd_warp_sync_functions.h | 469 ++++++++++++++++++ hipamd/src/hiprtc/CMakeLists.txt | 8 +- 2 files changed, 473 insertions(+), 4 deletions(-) diff --git a/hipamd/include/hip/amd_detail/amd_warp_sync_functions.h b/hipamd/include/hip/amd_detail/amd_warp_sync_functions.h index bd78b47aa6..80484b884e 100644 --- a/hipamd/include/hip/amd_detail/amd_warp_sync_functions.h +++ b/hipamd/include/hip/amd_detail/amd_warp_sync_functions.h @@ -32,8 +32,51 @@ THE SOFTWARE. #if !defined(__HIPCC_RTC__) #include "amd_warp_functions.h" #include "hip_assert.h" +#include "amd_hip_fp16.h" +#include +#include +#include #endif +extern "C" __device__ __attribute__((const)) int __ockl_wfred_add_i32(int); +extern "C" __device__ __attribute__((const)) unsigned int __ockl_wfred_add_u32(unsigned int); +extern "C" __device__ __attribute__((const)) long long __ockl_wfred_add_i64(long long); +extern "C" __device__ __attribute__((const)) unsigned long long __ockl_wfred_add_u64(unsigned long long); +extern "C" __device__ __attribute__((const)) __half __ockl_wfred_add_f16(__half); +extern "C" __device__ __attribute__((const)) float __ockl_wfred_add_f32(float); +extern "C" __device__ __attribute__((const)) double __ockl_wfred_add_f64(double); + +extern "C" __device__ __attribute__((const)) int __ockl_wfred_min_i32(int); +extern "C" __device__ __attribute__((const)) unsigned int __ockl_wfred_min_u32(unsigned int); +extern "C" __device__ __attribute__((const)) long long __ockl_wfred_min_i64(long long); +extern "C" __device__ __attribute__((const)) unsigned long long __ockl_wfred_min_u64(unsigned long long); +extern "C" __device__ __attribute__((const)) __half __ockl_wfred_min_f16(__half); +extern "C" __device__ __attribute__((const)) float __ockl_wfred_min_f32(float); +extern "C" __device__ __attribute__((const)) double __ockl_wfred_min_f64(double); + +extern "C" __device__ __attribute__((const)) int __ockl_wfred_max_i32(int); +extern "C" __device__ __attribute__((const)) unsigned int __ockl_wfred_max_u32(unsigned int); +extern "C" __device__ __attribute__((const)) long long __ockl_wfred_max_i64(long long); +extern "C" __device__ __attribute__((const)) unsigned long long __ockl_wfred_max_u64(unsigned long long); +extern "C" __device__ __attribute__((const)) __half __ockl_wfred_max_f16(__half); +extern "C" __device__ __attribute__((const)) float __ockl_wfred_max_f32(float); +extern "C" __device__ __attribute__((const)) double __ockl_wfred_max_f64(double); + +extern "C" __device__ __attribute__((const)) int __ockl_wfred_and_i32(int); +extern "C" __device__ __attribute__((const)) unsigned int __ockl_wfred_and_u32(unsigned int); +extern "C" __device__ __attribute__((const)) long long __ockl_wfred_and_i64(long long); +extern "C" __device__ __attribute__((const)) unsigned long long __ockl_wfred_and_u64(unsigned long long); + +extern "C" __device__ __attribute__((const)) int __ockl_wfred_or_i32(int); +extern "C" __device__ __attribute__((const)) unsigned int __ockl_wfred_or_u32(unsigned int); +extern "C" __device__ __attribute__((const)) long long __ockl_wfred_or_i64(long long); +extern "C" __device__ __attribute__((const)) unsigned long long __ockl_wfred_or_u64(unsigned long long); + +extern "C" __device__ __attribute__((const)) int __ockl_wfred_xor_i32(int); +extern "C" __device__ __attribute__((const)) unsigned int __ockl_wfred_xor_u32(unsigned int); +extern "C" __device__ __attribute__((const)) long long __ockl_wfred_xor_i64(long long); +extern "C" __device__ __attribute__((const)) unsigned long long __ockl_wfred_xor_u64(unsigned long long); + template __device__ inline T __hip_readfirstlane(T val) { @@ -279,6 +322,432 @@ T __shfl_xor_sync(MaskT mask, T var, int laneMask, return __shfl_xor(var, laneMask, width); } +template +__device__ inline T __reduce_op_sync(MaskT mask, T val, BinaryOp op, WfReduce wfReduce) +{ + static constexpr auto kMaskNumBits = sizeof(MaskT) * 8; + 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_adjust_mask_for_wave32(mask); + unsigned int laneId; + unsigned int maskIdx; + // next bit to aggregate with + int nextBit; + + // if doing the binary reduction tree, this will increase by two in every iteration + int modulo = 1; + int leadingZeroes = __clzll(mask); + int firstLane; + int lastLane = kMaskNumBits - leadingZeroes - 1; + int maskNumBits; + int numIterations; + // int[2] is used when T is 64-bit wide + typename __hip_internal::conditional::type result, permuteResult; + auto backwardPermute = [](auto index, auto val) { + if constexpr (std::is_integral::value || std::is_same::value) + return __hip_ds_bpermute(index, val); + else + return __hip_ds_bpermutef(index, val); + }; + + __hip_check_mask(mask); + maskNumBits = __popcll(mask); + +#ifdef __OPTIMIZE__ // at the time of this writing the ockl wfred functions do not compile when using -O0 + if (maskNumBits == lastLane + 1) + // this means the mask "does not have holes", and starts from 0; we can use a specific intrinsic + // to calculate the aggregated result + return wfReduce(val); +#endif + + firstLane = __builtin_ctzll(mask); + laneId = __ockl_lane_u32(); + nextBit = laneId; + // the number of iterations needs to be at least log2(number of bits on) + numIterations = sizeof(int) * 8 - __clz(maskNumBits); + + if (!(maskNumBits & (maskNumBits - 1))) + // the number of bits in the mask is a power of 2 + numIterations -= 1; + + mask >>= laneId; + mask >>= 1ul; + + if constexpr(sizeof(T) == 4 || sizeof(T) == 2) + result = val; + else + __builtin_memcpy(&result, &val, sizeof(T)); + + maskIdx = __ockl_activelane_u32(); + + // add the values from the lanes using a reduction tree (first the threads with even-numbered + // lanes, then multiples of 4, then 8, ... + while (numIterations) { + int offset = modulo >> 1; + int increment = modulo - offset; + int nextPos = maskIdx + offset + increment; + bool insideLanes = nextPos < maskNumBits; + + if (insideLanes) { + int next; + + // find the position to aggregate with; although we could just call fns64() that will probably + // be very slow when called multiple times in this for loop; this is equivalent + for (int i = 0; i < increment; i++) { + next = __builtin_ctzll(mask) + 1; + mask >>= next; + nextBit += next; + } + } + + if constexpr (std::is_same::value) { + union { int i; __half f; } tmp; + + tmp.f = result; + tmp.i = __hip_ds_bpermute(nextBit << 2, tmp.i); + permuteResult = tmp.f; + } else if constexpr (sizeof(T) == 4) + permuteResult = backwardPermute(nextBit << 2, result); + else { + // ds_bpermute only deals with 32-bit sizes, so for 64-bit types + // we need to call the permute twice for each half + permuteResult[0] = backwardPermute(nextBit << 2, result[0]); + permuteResult[1] = backwardPermute(nextBit << 2, result[1]); + } + + if (insideLanes) { + if constexpr (sizeof(T) == 4 || sizeof(T) == 2) + result = op(result, permuteResult); + else { + T tmp; + unsigned long long rhs = (static_cast(permuteResult[1]) << 32) | permuteResult[0]; + + __builtin_memcpy(&tmp, &result, sizeof(T)); + tmp = op(tmp, *reinterpret_cast(&rhs)); + __builtin_memcpy(&result, &tmp, sizeof(T)); + } + } + + modulo <<= 1; + numIterations--; + } + + if constexpr (std::is_same::value) { + union { int i; __half f; } tmp; + tmp.f = result; + tmp.i = __hip_ds_bpermute(firstLane << 2, tmp.i); + return tmp.f; + } else if constexpr (sizeof(T) == 4) + return backwardPermute(firstLane << 2, result); + else { + auto tmp = (static_cast(backwardPermute(firstLane << 2, result[1])) << 32) | + static_cast(backwardPermute(firstLane << 2, result[0])); + return *reinterpret_cast(&tmp); + } +} + +template +__device__ inline int __reduce_add_sync(MaskT mask, int val) +{ + // although C++ has std::plus and other functors, we do not use them because + // the they are header and they were causing problem with hipRTC + // at this time + auto op = [](decltype(val)& a, decltype(val)& b) { return a + b; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_add_i32(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline unsigned int __reduce_add_sync(MaskT mask, unsigned int val) +{ + auto op = [](decltype(val)& a, decltype(val)& b) { return a + b; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_add_u32(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline long long __reduce_add_sync(MaskT mask, long long val) +{ + auto op = [](decltype(val)& a, decltype(val)& b) { return a + b; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_add_i64(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline unsigned long long __reduce_add_sync(MaskT mask, unsigned long long val) +{ + auto op = [](decltype(val)& a, decltype(val)& b) { return a + b; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_add_u64(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline __half __reduce_add_sync(MaskT mask, __half val) +{ + auto op = [](decltype(val)& a, decltype(val)& b) { return a + b; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_add_f16(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline float __reduce_add_sync(MaskT mask, float val) +{ + auto op = [](decltype(val)& a, decltype(val)& b) { return a + b; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_add_f32(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline double __reduce_add_sync(MaskT mask, double val) +{ + auto op = [](decltype(val)& a, decltype(val)& b) { return a + b; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_add_f64(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline int __reduce_min_sync(MaskT mask, int val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return rhs < lhs? rhs : lhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_min_i32(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline unsigned int __reduce_min_sync(MaskT mask, unsigned int val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return rhs < lhs? rhs : lhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_min_u32(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline long long __reduce_min_sync(MaskT mask, long long val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return rhs < lhs? rhs : lhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_min_i64(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline unsigned long long __reduce_min_sync(MaskT mask, unsigned long long val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return rhs < lhs? rhs : lhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_min_u64(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline __half __reduce_min_sync(MaskT mask, __half val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return rhs < lhs? rhs : lhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_min_f16(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline float __reduce_min_sync(MaskT mask, float val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return rhs < lhs? rhs : lhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_min_f32(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline double __reduce_min_sync(MaskT mask, double val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return rhs < lhs? rhs : lhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_min_f64(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline int __reduce_max_sync(MaskT mask, int val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return lhs < rhs? rhs : lhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_max_i32(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline unsigned int __reduce_max_sync(MaskT mask, unsigned int val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return lhs < rhs? rhs : lhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_max_u32(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline long long __reduce_max_sync(MaskT mask, long long val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return lhs < rhs? rhs : lhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_max_i64(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline unsigned long long __reduce_max_sync(MaskT mask, unsigned long long val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return lhs < rhs? rhs : lhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_max_u64(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline __half __reduce_max_sync(MaskT mask, __half val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return lhs < rhs? rhs : lhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_max_f16(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline float __reduce_max_sync(MaskT mask, float val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return lhs < rhs? rhs : lhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_max_f32(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline double __reduce_max_sync(MaskT mask, double val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return lhs < rhs? rhs : lhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_max_f64(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline int __reduce_and_sync(MaskT mask, int val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return lhs && rhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_and_i32(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline unsigned int __reduce_and_sync(MaskT mask, unsigned int val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return lhs && rhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_and_u32(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline long long __reduce_and_sync(MaskT mask, long long val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return lhs && rhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_and_i64(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline unsigned long long __reduce_and_sync(MaskT mask, unsigned long long val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return lhs && rhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_and_u64(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline int __reduce_or_sync(MaskT mask, int val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return lhs || rhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_or_i32(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline unsigned int __reduce_or_sync(MaskT mask, unsigned int val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return lhs || rhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_or_u32(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline long long __reduce_or_sync(MaskT mask, long long val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return lhs || rhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_or_i64(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline unsigned long long __reduce_or_sync(MaskT mask, unsigned long long val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return lhs || rhs; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_or_u64(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline int __reduce_xor_sync(MaskT mask, int val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return (!lhs) != (!rhs) == 1; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_xor_i32(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline unsigned int __reduce_xor_sync(MaskT mask, unsigned int val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return (!lhs) != (!rhs) == 1; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_xor_u32(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline long long __reduce_xor_sync(MaskT mask, long long val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return (!lhs) != (!rhs) == 1; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_xor_i64(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + +template +__device__ inline unsigned long long __reduce_xor_sync(MaskT mask, unsigned long long val) +{ + auto op = [](decltype(val) lhs, decltype(val) rhs) { return (!lhs) != (!rhs)== 1; }; + auto wfReduce = [](decltype(val) v) { return __ockl_wfred_xor_u64(v); }; + + return __reduce_op_sync(mask, val, op, wfReduce); +} + #undef __hip_do_sync #undef __hip_check_mask #undef __hip_adjust_mask_for_wave32 diff --git a/hipamd/src/hiprtc/CMakeLists.txt b/hipamd/src/hiprtc/CMakeLists.txt index ef473f6237..f41573d66a 100644 --- a/hipamd/src/hiprtc/CMakeLists.txt +++ b/hipamd/src/hiprtc/CMakeLists.txt @@ -175,15 +175,15 @@ ${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_warp_functions.h ${PROJECT_SOURCE_DIR}/include/hip/amd_detail/amd_device_functions.h +${PROJECT_SOURCE_DIR}/include/hip/amd_detail/math_fwd.h +${PROJECT_SOURCE_DIR}/include/hip/amd_detail/hip_fp16_math_fwd.h +${PROJECT_SOURCE_DIR}/include/hip/amd_detail/amd_hip_fp16.h +${PROJECT_SOURCE_DIR}/include/hip/amd_detail/amd_math_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 ${PROJECT_SOURCE_DIR}/include/hip/amd_detail/amd_hip_atomic.h -${PROJECT_SOURCE_DIR}/include/hip/amd_detail/math_fwd.h -${PROJECT_SOURCE_DIR}/include/hip/amd_detail/hip_fp16_math_fwd.h -${PROJECT_SOURCE_DIR}/include/hip/amd_detail/amd_hip_fp16.h -${PROJECT_SOURCE_DIR}/include/hip/amd_detail/amd_math_functions.h ) # Generate required HIPRTC files.