SWDEV-420237 - Add __reduce_add_sync()

Change-Id: Ic8e4fab6b7aeb879d40b2c1419b30d1355a2bbdc

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 <functional>
+#include <algorithm>
+#include <cstring>
 #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 <typename T>
 __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 <typename MaskT, typename T, typename BinaryOp, typename WfReduce>
+__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<MaskT>::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<sizeof(T) == 4 || sizeof(T) == 2, T, unsigned int[2]>::type result, permuteResult;
+  auto backwardPermute = [](auto index, auto val) {
+    if constexpr (std::is_integral<T>::value || std::is_same<T, double>::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<T, __half>::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<unsigned long long>(permuteResult[1]) << 32) | permuteResult[0];
+
+        __builtin_memcpy(&tmp, &result, sizeof(T));
+        tmp = op(tmp, *reinterpret_cast<T*>(&rhs));
+        __builtin_memcpy(&result, &tmp, sizeof(T));
+      }
+    }
+
+    modulo <<= 1;
+    numIterations--;
+  }
+
+  if constexpr (std::is_same<T, __half>::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<unsigned long long>(backwardPermute(firstLane << 2, result[1])) << 32) |
+                static_cast<unsigned int>(backwardPermute(firstLane << 2, result[0]));
+    return *reinterpret_cast<T*>(&tmp);
+  }
+}
+
+template <typename MaskT>
+__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 <functional> 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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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 <typename MaskT>
+__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

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.