GPUCommonAlgorithm.h

Go to the documentation of this file.

// Copyright 2019-2020 CERN and copyright holders of ALICE O2.
// See https://alice-o2.web.cern.ch/copyright for details of the copyright holders.
// All rights not expressly granted are reserved.
//
// This software is distributed under the terms of the GNU General Public
// License v3 (GPL Version 3), copied verbatim in the file "COPYING".
//
// In applying this license CERN does not waive the privileges and immunities
// granted to it by virtue of its status as an Intergovernmental Organization
// or submit itself to any jurisdiction.
 
 
#ifndef GPUCOMMONALGORITHM_H
#define GPUCOMMONALGORITHM_H
 
#include "GPUCommonDef.h"
 
#if !defined(GPUCA_GPUCODE) || defined(__CPUDEVICE__) // Could also enable custom search on the CPU, but it is not always faster, so we stick to std::sort
#ifndef GPUCA_GPUCODE_DEVICE
#include <algorithm>
#endif
#define GPUCA_ALGORITHM_STD
#endif
 
// ----------------------------- SORTING -----------------------------
 
namespace o2::gpu
{
class GPUCommonAlgorithm
{
 public:
  template <class T>
  GPUd() static void sort(T* begin, T* end);
  template <class T>
  GPUd() static void sortInBlock(T* begin, T* end GPUCA_THREAD_INFO_DECL);
  template <class T>
  GPUd() static void sortDeviceDynamic(T* begin, T* end);
  template <class T, class S>
  GPUd() static void sort(T* begin, T* end, const S& comp);
  template <class T, class S>
  GPUd() static void sortInBlock(T* begin, T* end, const S& comp GPUCA_THREAD_INFO_DECL);
  template <class T, class S>
  GPUd() static void sortDeviceDynamic(T* begin, T* end, const S& comp);
#if __cplusplus >= 202002L // FIXME: auto parameter not supported in C++17
  template <class T, class S>
  GPUh() static void sortOnDevice(auto* rec, int32_t stream, T* begin, size_t N, const S& comp);
#endif
  template <class T>
  GPUd() static void swap(T& a, T& b);
  template <class I, class P>
  GPUhd() static I partitionPoint(I first, I count, const P& pred);
 
 private:
  template <typename I>
  GPUd() static void QuickSort(I f, I l) GPUnoexcept(); // Quicksort implementation
 
  template <typename I, typename Cmp>
  GPUd() static void QuickSort(I f, I l, Cmp cmp) GPUnoexcept(); // Quicksort implementation
 
  template <typename IndexType, typename I, typename Cmp>
  GPUd() static void QuickSortImpl(I f, I l, Cmp cmp) GPUnoexcept(); // Quicksort implementation with index type IndexType, which must be able to hold the number of elements
 
  template <typename I, typename Cmp>
  GPUd() static void InsertionSort(I f, I l, Cmp cmp) GPUnoexcept(); // Insertionsort implementation
 
  template <typename I, typename Cmp>
  GPUd() static I MedianOf3Select(I f, I l, Cmp cmp) GPUnoexcept(); // Helper for Quicksort implementation
 
  template <typename I, typename T, typename Cmp>
  GPUd() static I UnguardedPartition(I f, I l, T piv, Cmp cmp) GPUnoexcept(); // Helper for Quicksort implementation
 
  template <typename I>
  GPUd() static void IterSwap(I a, I b) GPUnoexcept(); // Helper
};
 
#ifndef GPUCA_ALGORITHM_STD
template <typename I>
GPUdi() void GPUCommonAlgorithm::IterSwap(I a, I b) GPUnoexcept()
{
  auto tmp = *a;
  *a = *b;
  *b = tmp;
}
 
template <typename I, typename Cmp>
GPUdi() void GPUCommonAlgorithm::InsertionSort(I f, I l, Cmp cmp) GPUnoexcept()
{
  auto it0{f};
  while (it0 != l) {
    auto tmp{*it0};
 
    auto it1{it0};
    while (it1 != f && cmp(tmp, it1[-1])) {
      it1[0] = it1[-1];
      --it1;
    }
    it1[0] = tmp;
 
    ++it0;
  }
}
 
template <typename I, typename Cmp>
GPUdi() I GPUCommonAlgorithm::MedianOf3Select(I f, I l, Cmp cmp) GPUnoexcept()
{
  auto m = f + (l - f) / 2;
 
  --l;
 
  if (cmp(*f, *m)) {
    if (cmp(*m, *l)) {
      return m;
    } else if (cmp(*f, *l)) {
      return l;
    } else {
      return f;
    }
  } else if (cmp(*f, *l)) {
    return f;
  } else if (cmp(*m, *l)) {
    return l;
  } else {
    return m;
  }
}
 
template <typename I, typename T, typename Cmp>
GPUdi() I GPUCommonAlgorithm::UnguardedPartition(I f, I l, T piv, Cmp cmp) GPUnoexcept()
{
  do {
    while (cmp(*f, piv)) {
      ++f;
    }
    --l;
    while (cmp(piv, *l)) {
      --l;
    }
 
    if (l <= f) {
      return f;
    }
    IterSwap(f, l);
    ++f;
  } while (true);
}
 
template <typename I, typename Cmp>
GPUdi() void GPUCommonAlgorithm::QuickSort(I f, I l, Cmp cmp) GPUnoexcept()
{
  if (l - f <= 0xFFFF) {
    QuickSortImpl<uint16_t>(f, l, cmp);
  } else {
    QuickSortImpl<uint32_t>(f, l, cmp);
  }
}
 
template <typename IndexType, typename I, typename Cmp>
GPUdi() void GPUCommonAlgorithm::QuickSortImpl(I f, I l, Cmp cmp) GPUnoexcept()
{
  if (f == l) {
    return;
  }
 
  struct pair {
    IndexType first;
    IndexType second;
  };
 
  struct Stack {
    pair data[sizeof(IndexType) * 8 - 4]; // The smaller partition is processed first, so the depth is bounded by log2(n / cutoff) + 1, cutoff = 50 > 2^5
    uint8_t n{0};
 
    GPUd() void emplace(IndexType x, IndexType y)
    {
      data[n++] = {x, y};
    }
    GPUd() bool empty() const { return n == 0; }
    GPUd() pair& top() { return data[n - 1]; }
    GPUd() void pop() { --n; }
  };
 
  Stack s;
  s.emplace(0, l - f);
  while (!s.empty()) {
    const auto it0 = f + s.top().first;
    const auto it1 = f + s.top().second;
    s.pop();
 
    const auto piv = *MedianOf3Select(it0, it1, cmp);
    const auto pp = UnguardedPartition(it0, it1, piv, cmp);
 
    constexpr auto cutoff = 50u;
    const auto lsz = pp - it0;
    const auto rsz = it1 - pp;
    if (lsz < rsz) {
      if (rsz > cutoff) {
        s.emplace(pp - f, it1 - f);
      }
      if (lsz > cutoff) {
        s.emplace(it0 - f, pp - f);
      }
    } else {
      if (lsz > cutoff) {
        s.emplace(it0 - f, pp - f);
      }
      if (rsz > cutoff) {
        s.emplace(pp - f, it1 - f);
      }
    }
  }
  InsertionSort(f, l, cmp);
}
 
template <typename I>
GPUdi() void GPUCommonAlgorithm::QuickSort(I f, I l) GPUnoexcept()
{
  QuickSort(f, l, [](auto&& x, auto&& y) { return x < y; });
}
#endif
 
typedef GPUCommonAlgorithm CAAlgo;
 
} // namespace o2::gpu
 
#if (((defined(__CUDACC__) && !defined(__clang__)) || defined(__HIPCC__))) && !defined(GPUCA_GPUCODE_HOSTONLY)
 
#include "GPUCommonAlgorithmThrust.h"
 
#else
 
namespace o2::gpu
{
 
template <class T>
GPUdi() void GPUCommonAlgorithm::sortDeviceDynamic(T* begin, T* end)
{
#ifndef GPUCA_GPUCODE
  GPUCommonAlgorithm::sort(begin, end);
#else
  GPUCommonAlgorithm::sortDeviceDynamic(begin, end, [](auto&& x, auto&& y) { return x < y; });
#endif
}
 
template <class T, class S>
GPUdi() void GPUCommonAlgorithm::sortDeviceDynamic(T* begin, T* end, const S& comp)
{
  GPUCommonAlgorithm::sort(begin, end, comp);
}
 
} // namespace o2::gpu
 
#endif // THRUST
// sort and sortInBlock below are not taken from Thrust, since our implementations are faster
 
namespace o2::gpu
{
 
template <class T>
GPUdi() void GPUCommonAlgorithm::sort(T* begin, T* end)
{
#ifdef GPUCA_ALGORITHM_STD
  std::sort(begin, end);
#else
  QuickSort(begin, end, [](auto&& x, auto&& y) { return x < y; });
#endif
}
 
template <class T, class S>
GPUdi() void GPUCommonAlgorithm::sort(T* begin, T* end, const S& comp)
{
#ifdef GPUCA_ALGORITHM_STD
  std::sort(begin, end, comp);
#else
  QuickSort(begin, end, comp);
#endif
}
 
template <class T>
GPUdi() void GPUCommonAlgorithm::sortInBlock(T* begin, T* end GPUCA_THREAD_INFO_DECL)
{
#ifndef GPUCA_GPUCODE
  GPUCommonAlgorithm::sort(begin, end);
#else
  GPUCommonAlgorithm::sortInBlock(begin, end, [](auto&& x, auto&& y) { return x < y; } GPUCA_THREAD_INFO_PROVIDE);
#endif
}
 
template <class T, class S>
GPUdi() void GPUCommonAlgorithm::sortInBlock(T* begin, T* end, const S& comp GPUCA_THREAD_INFO_DECL)
{
#if !defined(GPUCA_GPUCODE) || defined(GPUCA_DEVICE_SINGLE_THREAD)
  GPUCommonAlgorithm::sort(begin, end, comp);
#elif defined(GPUCA_DETERMINISTIC_MODE) // Cannot use GPUCA_DETERMINISTIC_CODE macro here, which is enforced in TPC compression
  if (get_local_id(0) == 0) {
    GPUCommonAlgorithm::sort(begin, end, comp);
  }
  GPUbarrier();
#else
  int32_t n = end - begin;
  for (int32_t i = 0; i < n; i++) {
    for (int32_t tIdx = get_local_id(0); tIdx < n; tIdx += get_local_size(0)) {
      int32_t offset = i % 2;
      int32_t curPos = 2 * tIdx + offset;
      int32_t nextPos = curPos + 1;
 
      if (nextPos < n) {
        if (!comp(begin[curPos], begin[nextPos])) {
          IterSwap(&begin[curPos], &begin[nextPos]);
        }
      }
    }
    GPUbarrier();
  }
#endif
}
 
#ifdef GPUCA_GPUCODE
template <class T>
GPUdi() void GPUCommonAlgorithm::swap(T& a, T& b)
{
  auto tmp = a;
  a = b;
  b = tmp;
}
#else
template <class T>
GPUdi() void GPUCommonAlgorithm::swap(T& a, T& b)
{
  std::swap(a, b);
}
#endif
 
// std::partition_point over the indices [first, first + count): the first index for which pred is false, pred being true on a prefix.
// A lower bound of value in a sorted key(i) is pred(i) = key(i) < value, an upper bound pred(i) = !(value < key(i)).
template <class I, class P>
GPUhdi() I GPUCommonAlgorithm::partitionPoint(I first, I count, const P& pred)
{
  while (count > 0) {
    const I step = count / 2;
    const I middle = first + step;
    if (pred(middle)) {
      first = middle + 1;
      count -= step + 1;
    } else {
      count = step;
    }
  }
  return first;
}
 
} // namespace o2::gpu
 
// ----------------------------- WORK GROUP FUNCTIONS -----------------------------
 
#ifdef __OPENCL__
// Nothing to do, work_group functions available
#pragma OPENCL EXTENSION cl_khr_subgroups : enable
 
template <class T>
GPUdi() T warp_scan_inclusive_add_FUNC(T v)
{
  return sub_group_scan_inclusive_add(v);
}
template <> // FIXME: It seems OpenCL does not support 8 and 16 bit subgroup operations
GPUdi() uint8_t warp_scan_inclusive_add_FUNC<uint8_t>(uint8_t v)
{
  return sub_group_scan_inclusive_add((uint32_t)v);
}
template <class T>
GPUdi() T warp_broadcast_FUNC(T v, int32_t i)
{
  return sub_group_broadcast(v, i);
}
template <>
GPUdi() uint8_t warp_broadcast_FUNC<uint8_t>(uint8_t v, int32_t i)
{
  return sub_group_broadcast((uint32_t)v, i);
}
 
#define warp_scan_inclusive_add(v) warp_scan_inclusive_add_FUNC(v)
#define warp_broadcast(v, i) warp_broadcast_FUNC(v, i)
 
[[nodiscard]] GPUdi() int32_t work_group_count(bool pred)
{
  return work_group_reduce_add((int32_t)pred);
}
 
#elif (defined(__CUDACC__) || defined(__HIPCC__))
// CUDA and HIP work the same way using cub, need just different header
 
#if !defined(GPUCA_GPUCODE_COMPILEKERNELS) && !defined(GPUCA_GPUCODE_HOSTONLY)
#if defined(__CUDACC__)
#include <cub/cub.cuh>
#elif defined(__HIPCC__)
#include <hipcub/hipcub.hpp>
#endif
#endif
 
#define work_group_scan_inclusive_add(v) work_group_scan_inclusive_add_FUNC(v, smem)
template <class T, class S>
GPUdi() T work_group_scan_inclusive_add_FUNC(T v, S& smem)
{
  typename S::BlockScan(smem.cubTmpMem).InclusiveSum(v, v);
  __syncthreads();
  return v;
}
 
#define work_group_broadcast(v, i) work_group_broadcast_FUNC(v, i, smem)
template <class T, class S>
GPUdi() T work_group_broadcast_FUNC(T v, int32_t i, S& smem)
{
  if ((int32_t)threadIdx.x == i) {
    smem.tmpBroadcast = v;
  }
  __syncthreads();
  T retVal = smem.tmpBroadcast;
  __syncthreads();
  return retVal;
}
 
#define work_group_reduce_add(v) work_group_reduce_add_FUNC(v, smem)
template <class T, class S>
GPUdi() T work_group_reduce_add_FUNC(T v, S& smem)
{
  v = typename S::BlockReduce(smem.cubReduceTmpMem).Sum(v);
  __syncthreads();
  v = work_group_broadcast(v, 0);
  return v;
}
 
#define warp_scan_inclusive_add(v) warp_scan_inclusive_add_FUNC(v, smem)
template <class T, class S>
GPUdi() T warp_scan_inclusive_add_FUNC(T v, S& smem)
{
  typename S::WarpScan(smem.cubWarpTmpMem).InclusiveSum(v, v);
  return v;
}
 
#define warp_broadcast(v, i) warp_broadcast_FUNC(v, i)
template <class T>
GPUdi() T warp_broadcast_FUNC(T v, int32_t i)
{
#ifdef __CUDACC__
  return __shfl_sync(0xFFFFFFFF, v, i);
#else // HIP
  return __shfl(v, i);
#endif
}
 
template <class T>
GPUdi() T warp_shuffle_xor(T v, int32_t laneMask)
{
#ifdef __CUDACC__
  return __shfl_xor_sync(0xFFFFFFFF, v, laneMask);
#else // HIP
  return __shfl_xor(v, laneMask);
#endif
}
 
// Butterfly all-reduce over groups of Width lanes: every lane gets the result.
// Width is a power of two and must not exceed the hardware warp size.
template <int32_t Width, class T>
GPUdi() T warp_reduce_add(T v)
{
  for (int32_t m = Width / 2; m > 0; m >>= 1) {
    v += warp_shuffle_xor(v, m);
  }
  return v;
}
 
template <int32_t Width, class T>
GPUdi() T warp_reduce_min(T v)
{
  for (int32_t m = Width / 2; m > 0; m >>= 1) {
    const T o = warp_shuffle_xor(v, m);
    v = o < v ? o : v;
  }
  return v;
}
 
template <int32_t Width, class T>
GPUdi() T warp_reduce_max(T v)
{
  for (int32_t m = Width / 2; m > 0; m >>= 1) {
    const T o = warp_shuffle_xor(v, m);
    v = o > v ? o : v;
  }
  return v;
}
 
[[nodiscard]] GPUdi() bool work_group_any(bool pred)
{
  return __syncthreads_or(pred);
}
 
[[nodiscard]] GPUdi() uint32_t work_group_count(bool pred)
{
  return __syncthreads_count(pred);
}
 
#elif defined(__METAL__)
// MSL has SIMD-group collectives only, so the work-group ones scan within each
// SIMD group and then across the group totals, which go through smem as with
// cub. The thread index is a kernel parameter in MSL, so like smem it is taken
// from the caller's scope. A 1D threadgroup is split into SIMD groups of
// consecutive threads, of GPUCA_WARP_SIZE each. The SIMD scans take 32-bit
// lanes; the narrower types are counts and are widened. Sums are returned
// promoted, as v + 0 would be, since the callers pass !!pred: a bool would
// truncate them to 0 / 1.
 
template <class T>
GPUdi() auto warp_scan_inclusive_add(T v) -> decltype(+v)
{
  return simd_prefix_inclusive_sum((int32_t)v);
}
 
template <class T>
GPUdi() T warp_broadcast(T v, int32_t i)
{
  return (T)simd_shuffle((int32_t)v, (ushort)i);
}
 
template <class T>
GPUdi() T warp_shuffle_xor(T v, int32_t laneMask)
{
  return simd_shuffle_xor(v, (ushort)laneMask);
}
 
template <int32_t Width, class T>
GPUdi() T warp_reduce_add(T v)
{
  for (int32_t m = Width / 2; m > 0; m >>= 1) {
    v += warp_shuffle_xor(v, m);
  }
  return v;
}
 
template <int32_t Width, class T>
GPUdi() T warp_reduce_min(T v)
{
  for (int32_t m = Width / 2; m > 0; m >>= 1) {
    const T o = warp_shuffle_xor(v, m);
    v = o < v ? o : v;
  }
  return v;
}
 
template <int32_t Width, class T>
GPUdi() T warp_reduce_max(T v)
{
  for (int32_t m = Width / 2; m > 0; m >>= 1) {
    const T o = warp_shuffle_xor(v, m);
    v = o > v ? o : v;
  }
  return v;
}
 
#define work_group_scan_inclusive_add(v) work_group_scan_inclusive_add_FUNC(v, smem, iThread, nThreads)
template <class T, class S>
GPUdi() auto work_group_scan_inclusive_add_FUNC(T v, S& smem, int32_t iThread, int32_t nThreads) -> decltype(+v)
{
  const int32_t iWarp = iThread / GPUCA_WARP_SIZE;
  const int32_t nWarps = (nThreads + GPUCA_WARP_SIZE - 1) / GPUCA_WARP_SIZE;
  int32_t x = simd_prefix_inclusive_sum((int32_t)v);
  if (iThread % GPUCA_WARP_SIZE == GPUCA_WARP_SIZE - 1 || iThread == nThreads - 1) {
    smem.warpPredicateSum[iWarp] = x;
  }
  GPUbarrier();
  if (iWarp == 0) {
    const int32_t s = simd_prefix_inclusive_sum(iThread < nWarps ? smem.warpPredicateSum[iThread] : 0);
    if (iThread < nWarps) {
      smem.warpPredicateSum[iThread] = s;
    }
  }
  GPUbarrier();
  if (iWarp > 0) {
    x += smem.warpPredicateSum[iWarp - 1];
  }
  GPUbarrier();
  return x;
}
 
#define work_group_broadcast(v, i) work_group_broadcast_FUNC(v, i, smem, iThread)
template <class T, class S>
GPUdi() T work_group_broadcast_FUNC(T v, int32_t i, S& smem, int32_t iThread)
{
  if (iThread == i) {
    smem.tmpBroadcast = v;
  }
  GPUbarrier();
  T retVal = smem.tmpBroadcast;
  GPUbarrier();
  return retVal;
}
 
#define work_group_reduce_add(v) work_group_reduce_add_FUNC(v, smem, iThread, nThreads)
template <class T, class S>
GPUdi() auto work_group_reduce_add_FUNC(T v, S& smem, int32_t iThread, int32_t nThreads) -> decltype(+v)
{
  return work_group_broadcast_FUNC(work_group_scan_inclusive_add_FUNC(v, smem, iThread, nThreads), nThreads - 1, smem, iThread);
}
 
#define work_group_count(pred) ((uint32_t)work_group_reduce_add_FUNC((int32_t)(pred), smem, iThread, nThreads))
#define work_group_any(pred) (work_group_count(pred) != 0)
 
#else
// Trivial implementation for the CPU
 
template <class T>
GPUdi() T work_group_scan_inclusive_add(T v)
{
  return v;
}
 
template <class T>
GPUdi() T work_group_reduce_add(T v)
{
  return v;
}
 
template <class T>
GPUdi() T work_group_broadcast(T v, int32_t i)
{
  return v;
}
 
template <class T>
GPUdi() T warp_scan_inclusive_add(T v)
{
  return v;
}
 
template <class T>
GPUdi() T warp_broadcast(T v, int32_t i)
{
  return v;
}
 
template <int32_t Width, class T>
GPUdi() T warp_reduce_add(T v)
{
  return v;
}
 
template <int32_t Width, class T>
GPUdi() T warp_reduce_min(T v)
{
  return v;
}
 
template <int32_t Width, class T>
GPUdi() T warp_reduce_max(T v)
{
  return v;
}
 
[[nodiscard]] GPUdi() bool work_group_any(bool pred)
{
  return pred;
}
 
[[nodiscard]] GPUdi() uint32_t work_group_count(bool pred)
{
  return pred;
}
 
#endif
 
#ifdef GPUCA_ALGORITHM_STD
#undef GPUCA_ALGORITHM_STD
#endif
 
#endif