EBGeometry_GPUAlgorithms.hpp

Go to the documentation of this file.

// SPDX-FileCopyrightText: 2026 Robert Marskar <robert.marskar@sintef.no>
//
// SPDX-License-Identifier: GPL-3.0-or-later
 
#ifndef EBGEOMETRY_GPUALGORITHMS_HPP
#define EBGEOMETRY_GPUALGORITHMS_HPP
 
// Std includes
#include <cmath>
#include <cstddef>
#include <cstdint>
#include <type_traits>
 
// Our includes
#include "EBGeometry_BoundingVolumes.hpp"
#include "EBGeometry_GPU.hpp"
#include "EBGeometry_GPURuntime.hpp"
#include "EBGeometry_GPUTransform.hpp"
#include "EBGeometry_Macros.hpp"
#include "EBGeometry_Math.hpp"
#include "EBGeometry_Vec.hpp"
 
namespace EBGeometry {
 
namespace GPU {
 
inline constexpr uint32_t ScanChunk = 256;
 
[[nodiscard]] EBGEOMETRY_HOST
inline uint32_t
exclusiveScanScratchSize(const uint32_t a_count) noexcept
{
  return (a_count + ScanChunk - 1) / ScanChunk + 1U;
}
 
inline constexpr uint32_t BoundingBoxChunk = 256;
 
[[nodiscard]] EBGEOMETRY_HOST
inline uint32_t
boundingBoxScratchSize(const uint32_t a_count) noexcept
{
  return 2U * Math::max(1U, (a_count + BoundingBoxChunk - 1) / BoundingBoxChunk);
}
 
namespace AlgorithmsDetail {
 
template <class T>
struct FromPointer
{
  const Vec3T<T>* m_points;
 
  [[nodiscard]] EBGEOMETRY_HOST_DEVICE
  Vec3T<T>
  operator()(const uint32_t a_i) const noexcept
  {
    return m_points[a_i];
  }
};
 
template <class T, class Source>
struct ChunkBox
{
  Source m_points;
 
  Vec3T<T>* m_corners;
 
  uint32_t m_count;
 
  EBGEOMETRY_HOST_DEVICE
  void
  operator()(const std::size_t a_c) const noexcept
  {
    const uint32_t begin = static_cast<uint32_t>(a_c) * BoundingBoxChunk;
    const uint32_t end   = Math::min(m_count, begin + BoundingBoxChunk);
 
    Vec3T<T> lo = +Vec3T<T>::max();
    Vec3T<T> hi = -Vec3T<T>::max();
 
    bool finite = true;
 
    for (uint32_t i = begin; i < end; i++) {
      const Vec3T<T> p = m_points(i);
 
      finite = finite && std::isfinite(p[0]) && std::isfinite(p[1]) && std::isfinite(p[2]);
 
      lo = min(lo, p);
      hi = max(hi, p);
    }
 
    // A chunk with a non-finite point marks its box with a lower corner of -infinity, which the merge keeps.
    if (!finite) {
      lo[0] = -Math::Limits<T>::infinity();
    }
 
    m_corners[2 * a_c]     = lo;
    m_corners[2 * a_c + 1] = hi;
  }
};
 
template <class T>
struct MergeChunks
{
  Vec3T<T>* m_corners;
 
  uint32_t m_chunks;
 
  EBGEOMETRY_HOST_DEVICE
  void
  operator()(std::size_t /*a_thread*/) const noexcept
  {
    Vec3T<T> lo = m_corners[0];
    Vec3T<T> hi = m_corners[1];
 
    for (uint32_t c = 1; c < m_chunks; c++) {
      lo = min(lo, m_corners[2 * c]);
      hi = max(hi, m_corners[2 * c + 1]);
    }
 
    m_corners[0] = lo;
    m_corners[1] = hi;
  }
};
 
template <class Key, class Less>
struct BitonicPass
{
  Key* m_keys;
 
  Less m_less;
 
  uint32_t m_k;
 
  uint32_t m_j;
 
  EBGEOMETRY_HOST_DEVICE
  void
  operator()(const std::size_t a_t) const noexcept
  {
    const auto     t = static_cast<uint32_t>(a_t);
    const uint32_t i = (t / m_j) * 2 * m_j + t % m_j;
    const uint32_t l = i + m_j;
 
    Key& a = m_keys[i];
    Key& b = m_keys[l];
 
    const bool ascending = (i & m_k) == 0;
 
    if (ascending ? m_less(b, a) : m_less(a, b)) {
      const Key swap = a;
 
      a = b;
      b = swap;
    }
  }
};
 
struct ChunkSums
{
  const uint32_t* m_values;
 
  uint32_t* m_sums;
 
  uint32_t m_count;
 
  EBGEOMETRY_HOST_DEVICE
  void
  operator()(const std::size_t a_c) const noexcept
  {
    const uint32_t begin = static_cast<uint32_t>(a_c) * ScanChunk;
    const uint32_t end   = Math::min(m_count, begin + ScanChunk);
 
    uint32_t sum = 0;
 
    for (uint32_t i = begin; i < end; i++) {
      sum += m_values[i];
    }
 
    m_sums[a_c] = sum;
  }
};
 
struct ScanSums
{
  uint32_t* m_sums;
 
  uint32_t m_chunks;
 
  EBGEOMETRY_HOST_DEVICE
  void
  operator()(std::size_t /*a_thread*/) const noexcept
  {
    uint32_t running = 0;
 
    for (uint32_t c = 0; c < m_chunks; c++) {
      const uint32_t sum = m_sums[c];
 
      m_sums[c] = running;
      running += sum;
    }
 
    m_sums[m_chunks] = running;
  }
};
 
struct ScanChunks
{
  uint32_t* m_values;
 
  const uint32_t* m_offsets;
 
  uint32_t m_count;
 
  EBGEOMETRY_HOST_DEVICE
  void
  operator()(const std::size_t a_c) const noexcept
  {
    const uint32_t begin = static_cast<uint32_t>(a_c) * ScanChunk;
    const uint32_t end   = Math::min(m_count, begin + ScanChunk);
 
    uint32_t running = m_offsets[a_c];
 
    for (uint32_t i = begin; i < end; i++) {
      const uint32_t value = m_values[i];
 
      m_values[i] = running;
      running += value;
    }
  }
};
 
} // namespace AlgorithmsDetail
 
struct Ascending
{
  template <class Key>
  [[nodiscard]] EBGEOMETRY_HOST_DEVICE
  bool
  operator()(const Key& a_lhs, const Key& a_rhs) const noexcept
  {
    return a_lhs < a_rhs;
  }
};
 
template <class T, class Source>
[[nodiscard]] EBGEOMETRY_HOST
inline BoundingVolumes::AABBT<T>
boundingBoxOf(const Source& a_points, const uint32_t a_count, Vec3T<T>* a_scratch)
{
  static_assert(std::is_trivially_copyable_v<Source>, "GPU::boundingBoxOf: the source must be trivially copyable");
 
  if (a_count == 0) {
    return BoundingVolumes::AABBT<T>();
  }
 
  const uint32_t chunks = boundingBoxScratchSize(a_count) / 2;
 
  parallelFor(chunks, AlgorithmsDetail::ChunkBox<T, Source>{a_points, a_scratch, a_count});
  parallelFor(1, AlgorithmsDetail::MergeChunks<T>{a_scratch, chunks});
 
  Vec3T<T> corners[2];
 
#if defined(EBGEOMETRY_CUDA) || defined(EBGEOMETRY_HIP)
  EBGEOMETRY_GPU_CHECK(GPU::memcpy(corners, a_scratch, sizeof(corners), GPU::MemcpyDeviceToHost));
#else
  corners[0] = a_scratch[0];
  corners[1] = a_scratch[1];
#endif
 
  return BoundingVolumes::AABBT<T>(corners[0], corners[1]);
}
 
template <class T>
[[nodiscard]] EBGEOMETRY_HOST
inline BoundingVolumes::AABBT<T>
boundingBox(const Vec3T<T>* a_points, const uint32_t a_count, Vec3T<T>* a_scratch)
{
  return boundingBoxOf(AlgorithmsDetail::FromPointer<T>{a_points}, a_count, a_scratch);
}
 
template <class Key, class Less = Ascending>
EBGEOMETRY_HOST
inline void
bitonicSort(Key* a_keys, const uint32_t a_count, const Less& a_less = Less())
{
  static_assert(std::is_trivially_copyable_v<Key>, "GPU::bitonicSort: the key type must be trivially copyable");
 
  EBGEOMETRY_REQUIRE(
    (a_count & (a_count - 1)) == 0, "GPU::bitonicSort: the key count must be a power of two (%u)", unsigned(a_count));
 
  for (uint32_t k = 2; k <= a_count && k != 0; k *= 2) {
    for (uint32_t j = k / 2; j > 0; j /= 2) {
      parallelFor(a_count / 2, AlgorithmsDetail::BitonicPass<Key, Less>{a_keys, a_less, k, j});
    }
  }
}
 
[[nodiscard]] EBGEOMETRY_HOST
inline uint32_t
exclusiveScan(uint32_t* a_values, const uint32_t a_count, uint32_t* a_scratch)
{
  if (a_count == 0) {
    return 0;
  }
 
  const uint32_t chunks = exclusiveScanScratchSize(a_count) - 1U;
 
  parallelFor(chunks, AlgorithmsDetail::ChunkSums{a_values, a_scratch, a_count});
  parallelFor(1, AlgorithmsDetail::ScanSums{a_scratch, chunks});
  parallelFor(chunks, AlgorithmsDetail::ScanChunks{a_values, a_scratch, a_count});
 
  uint32_t total = 0;
 
#if defined(EBGEOMETRY_CUDA) || defined(EBGEOMETRY_HIP)
  EBGEOMETRY_GPU_CHECK(GPU::memcpy(&total, a_scratch + chunks, sizeof(total), GPU::MemcpyDeviceToHost));
#else
  total = a_scratch[chunks];
#endif
 
  return total;
}
 
} // namespace GPU
 
} // namespace EBGeometry
 
#endif