AMReX_ParticleReduceSIMD.H

Go to the documentation of this file.

#ifndef AMREX_PARTICLEREDUCE_SIMD_H_
#define AMREX_PARTICLEREDUCE_SIMD_H_
#include <AMReX_Config.H>
 
#include <AMReX_ParticleReduce.H>  // provides amrex::SIMDindex via AMReX_Gpu.H
#include <AMReX_ReduceSIMD.H>
#include <AMReX_SIMD.H>
#include <AMReX_Tuple.H>
#include <AMReX_TypeTraits.H>
#include <AMReX_Vector.H>
 
#include <type_traits>
#include <utility>
 
namespace amrex {
 
namespace particle_detail {
 
    template <typename F>
    struct SIMDTupleWrapper
    {
        F m_f;
 
        template <typename PTD, typename SI>
        AMREX_GPU_HOST_DEVICE AMREX_FORCE_INLINE
        auto operator() (PTD const& ptd, SI const si) const noexcept
        {
            return amrex::makeTuple(m_f(ptd, si));
        }
    };
 
    template <typename F, typename PTD>
    struct SIMDKernelBinder
    {
        F m_f;
        PTD m_ptd;
 
        template <typename SI>
        AMREX_FORCE_INLINE
        auto operator() (SI const si) const noexcept
        {
            return m_f(m_ptd, si);
        }
    };
}
 
template <class RD,
          int WIDTH = static_cast<int>(simd::native_simd_size_particlereal),
          class PC, class F, typename... Ps>
  requires(IsParticleContainer<PC>::value)
typename RD::Type ParticleReduceSIMD(PC const &pc, int lev_min, int lev_max,
                                     F const &f, ReduceOps<Ps...> &reduce_ops) {
  // enforce the documented contract in every configuration (CPU with and
  // without SIMD support, GPU), not just where evalReduceSIMD is called
  static_assert(WIDTH >= 1, "ParticleReduceSIMD: WIDTH must be at least 1");
  static_assert((Reduce::detail::is_simd_reduce_op<Ps> && ...),
                "ParticleReduceSIMD supports only ReduceOpSum, ReduceOpMin and "
                "ReduceOpMax");
 
  RD reduce_data(reduce_ops);
  for (int lev = lev_min; lev <= lev_max; ++lev) {
    const auto &plev = pc.GetParticles(lev);
    Vector<std::pair<int, int>> grid_tile_ids;
    Vector<const typename PC::ParticleTileType *> ptile_ptrs;
    for (auto &kv : plev) {
      grid_tile_ids.push_back(kv.first);
      ptile_ptrs.push_back(&(kv.second));
    }
 
    // We avoid ParIter here to allow reductions after grids change but prior to
    // calling Redistribute()
#if !defined(AMREX_USE_GPU) && defined(AMREX_USE_OMP)
#pragma omp parallel for
#endif
        for (int pmap_it = 0; pmap_it < std::ssize(ptile_ptrs); ++pmap_it)
        {
            const auto& tile = plev.at(grid_tile_ids[pmap_it]);
            const auto np = tile.numParticles();
            const auto& ptd = tile.getConstParticleTileData();
#ifdef AMREX_USE_GPU
            reduce_ops.eval(np, reduce_data,
                            [=] AMREX_GPU_DEVICE (const int i) noexcept
                            {
                                return f(ptd, SIMDindex<1, int>{i});
                            });
#else
#ifdef AMREX_USE_SIMD
            constexpr int EffWidth = WIDTH;
#else
            constexpr int EffWidth = 1;  // explicit WIDTH > 1 still compiles, scalar
#endif
            using PTDType = std::decay_t<decltype(ptd)>;
            evalReduceSIMD<EffWidth>(np, reduce_data,
                            particle_detail::SIMDKernelBinder<F, PTDType>{f, ptd},
                            reduce_ops);
#endif
        }
    }
    return reduce_data.value(reduce_ops);
}
 
template <class RD,
          int WIDTH = static_cast<int>(simd::native_simd_size_particlereal),
          class PC, class F, typename... Ps>
  requires(IsParticleContainer<PC>::value)
typename RD::Type ParticleReduceSIMD(PC const &pc, int lev, F const &f,
                                     ReduceOps<Ps...> &reduce_ops) {
  return ParticleReduceSIMD<RD, WIDTH>(pc, lev, lev, f, reduce_ops);
}
 
template <class RD,
          int WIDTH = static_cast<int>(simd::native_simd_size_particlereal),
          class PC, class F, typename... Ps>
  requires(IsParticleContainer<PC>::value)
typename RD::Type ParticleReduceSIMD(PC const &pc, F const &f,
                                     ReduceOps<Ps...> &reduce_ops) {
  return ParticleReduceSIMD<RD, WIDTH>(pc, 0, pc.finestLevel(), f, reduce_ops);
}
 
template <int WIDTH = static_cast<int>(simd::native_simd_size_particlereal),
          class PC, class F>
requires (IsParticleContainer<PC>::value)
auto
ReduceSumSIMD (PC const& pc, int lev_min, int lev_max, F const& f)
{
    using ConstPTDType = typename PC::ConstPTDType;
    using T = std::decay_t<std::invoke_result_t<F, ConstPTDType const&, SIMDindex<1, int>>>;
    ReduceOps<ReduceOpSum> reduce_ops;
    auto r = ParticleReduceSIMD<ReduceData<T>, WIDTH>(
        pc, lev_min, lev_max, particle_detail::SIMDTupleWrapper<F>{f}, reduce_ops);
    return amrex::get<0>(r);
}
 
template <int WIDTH = static_cast<int>(simd::native_simd_size_particlereal),
          class PC, class F>
requires (IsParticleContainer<PC>::value)
auto
ReduceSumSIMD (PC const& pc, int lev, F const& f)
{
    return ReduceSumSIMD<WIDTH>(pc, lev, lev, f);
}
 
template <int WIDTH = static_cast<int>(simd::native_simd_size_particlereal),
          class PC, class F>
requires (IsParticleContainer<PC>::value)
auto
ReduceSumSIMD (PC const& pc, F const& f)
{
    return ReduceSumSIMD<WIDTH>(pc, 0, pc.finestLevel(), f);
}
 
template <int WIDTH = static_cast<int>(simd::native_simd_size_particlereal),
          class PC, class F>
requires (IsParticleContainer<PC>::value)
auto
ReduceMinSIMD (PC const& pc, int lev_min, int lev_max, F const& f)
{
    using ConstPTDType = typename PC::ConstPTDType;
    using T = std::decay_t<std::invoke_result_t<F, ConstPTDType const&, SIMDindex<1, int>>>;
    ReduceOps<ReduceOpMin> reduce_ops;
    auto r = ParticleReduceSIMD<ReduceData<T>, WIDTH>(
        pc, lev_min, lev_max, particle_detail::SIMDTupleWrapper<F>{f}, reduce_ops);
    return amrex::get<0>(r);
}
 
template <int WIDTH = static_cast<int>(simd::native_simd_size_particlereal),
          class PC, class F>
requires (IsParticleContainer<PC>::value)
auto
ReduceMinSIMD (PC const& pc, int lev, F const& f)
{
    return ReduceMinSIMD<WIDTH>(pc, lev, lev, f);
}
 
template <int WIDTH = static_cast<int>(simd::native_simd_size_particlereal),
          class PC, class F>
requires (IsParticleContainer<PC>::value)
auto
ReduceMinSIMD (PC const& pc, F const& f)
{
    return ReduceMinSIMD<WIDTH>(pc, 0, pc.finestLevel(), f);
}
 
template <int WIDTH = static_cast<int>(simd::native_simd_size_particlereal),
          class PC, class F>
requires (IsParticleContainer<PC>::value)
auto
ReduceMaxSIMD (PC const& pc, int lev_min, int lev_max, F const& f)
{
    using ConstPTDType = typename PC::ConstPTDType;
    using T = std::decay_t<std::invoke_result_t<F, ConstPTDType const&, SIMDindex<1, int>>>;
    ReduceOps<ReduceOpMax> reduce_ops;
    auto r = ParticleReduceSIMD<ReduceData<T>, WIDTH>(
        pc, lev_min, lev_max, particle_detail::SIMDTupleWrapper<F>{f}, reduce_ops);
    return amrex::get<0>(r);
}
 
template <int WIDTH = static_cast<int>(simd::native_simd_size_particlereal),
          class PC, class F>
requires (IsParticleContainer<PC>::value)
auto
ReduceMaxSIMD (PC const& pc, int lev, F const& f)
{
    return ReduceMaxSIMD<WIDTH>(pc, lev, lev, f);
}
 
template <int WIDTH = static_cast<int>(simd::native_simd_size_particlereal),
          class PC, class F>
requires (IsParticleContainer<PC>::value)
auto
ReduceMaxSIMD (PC const& pc, F const& f)
{
    return ReduceMaxSIMD<WIDTH>(pc, 0, pc.finestLevel(), f);
}
 
}
 
#endif