contribution_kernel.hpp

#ifndef __CONTRIBUTION_KERNEL_HPP__
#define __CONTRIBUTION_KERNEL_HPP__
#include "Kokkos_Assert.hpp"
#include "Kokkos_Macros.hpp"
#include "mc/events.hpp"
#include <Kokkos_Core.hpp>
#include <mc/particles_container.hpp>
#include <mc/traits.hpp>
 
template <ModelType M> struct ContributionFunctor
{
 
  struct Tag0D
  {
  };
  struct Tag3D
  {
  };
 
  using TeamPolicy = Kokkos::TeamPolicy<ComputeSpace>;
  using TeamMember = TeamPolicy::member_type;
  using float_type = float;
  using ScratchView
      = Kokkos::View<float_type*,
                     TeamPolicy::execution_space::scratch_memory_space>;
 
  ContributionFunctor(std::size_t particle_per_team,
                      MC::ContributionView contribution_scatter,
                      MC::ParticlesContainer<M> particles,
                      MC::pool_type rp,
                      MC::KernelConcentrationType _concentrations,
                      MC::EventContainer _event)
      : m_particle_per_team(particle_per_team),
        m_contribution_scatter(std::move(contribution_scatter)),
        m_particles(std::move(particles)), random_pool(rp),
        concentrations(std::move(_concentrations)), events(std::move(_event))
  {
  }
 
  size_t np{};
 
  void
  update(MC::ParticlesContainer<M> _particles, double _d_t)
 
  {
    this->d_t = _d_t;
    this->m_particles = std::move(_particles);
    np = m_particles.n_particles();
  }
 
  std::size_t m_particle_per_team;
 
  MC::ContributionView m_contribution_scatter;
  MC::ParticlesContainer<M> m_particles;
 
  MC::pool_type random_pool;
  MC::KernelConcentrationType concentrations;
  M::FloatType d_t;
  MC::EventContainer events;
 
  KOKKOS_INLINE_FUNCTION void
  division(const std::size_t idx) const
  {
 
    using mem_space = ComputeSpace::memory_space;
    if constexpr (AutoGenerated::FlagCompileTime::use_probe)
    {
      // Register probe here to sample BEFORE division and even if
      // division procedure fails to spawn new particle, the age is
      // still the same
      const auto _ = this->probes.template set<mem_space>(
          m_particles.ages(idx, 1)); // Skip error
    }
 
    if (!m_particles.handle_division(random_pool, idx)) [[unlikely]]
    {
      // reduce_val.waiting_allocation_particle += 1;
      // events.wrap_incr<MC::EventType::Overflow>();
      Kokkos::printf("[KERNEL] Division Overflow\r\n");
    }
 
    events.wrap_incr<MC::EventType::NewParticle>();
  }
 
  KOKKOS_INLINE_FUNCTION
  void
  operator()(Tag0D _tag, const TeamMember& team) const
  {
    (void)_tag;
    const std::size_t p0 = team.league_rank() * m_particle_per_team;
    const std::size_t n_particle = m_particles.n_particles();
    const auto& status = m_particles.status;
    const auto& contribs = m_particles.contribs;
    static const auto n_c = M::n_c;
 
    ScratchView scratch(team.team_scratch(0), n_c);
 
    Kokkos::parallel_for(Kokkos::TeamVectorRange(team, n_c),
                         [&](const std::size_t j)
                         { scratch(j) = float_type{ 0 }; });
 
    const auto upper_bound = ((p0 + m_particle_per_team) >= n_particle)
                                 ? n_particle - p0
                                 : m_particle_per_team;
 
    KOKKOS_ASSERT(upper_bound > 0 && upper_bound <= m_particle_per_team
                  && (p0 + upper_bound) <= n_particle);
 
    Kokkos::parallel_for(
        Kokkos::TeamThreadRange(team, 0, upper_bound),
        [&](std::size_t relative_index)
        {
          const std::size_t flatten_index = p0 + relative_index;
          KOKKOS_ASSERT(flatten_index < upper_bound);
          const bool active = status(flatten_index) == MC::Status::Idle;
 
          if (active)
          {
            const auto pos = m_particles.position(flatten_index);
            const auto weight = m_particles.get_weight(flatten_index);
            const auto new_status = M::update(random_pool,
                                              d_t,
                                              flatten_index,
                                              m_particles.model,
                                              contribs,
                                              pos,
                                              concentrations);
 
            for (std::size_t j = 0; j < n_c; ++j)
            {
              Kokkos::atomic_add(&scratch(j),
                                 weight * contribs(flatten_index, j));
            }
 
            if (new_status == MC::Status::Division)
            {
              division(flatten_index);
            }
            else
            {
              m_particles.ages(flatten_index, 1) += d_t;
            }
          }
        });
 
    team.team_barrier();
    const auto& cs = m_contribution_scatter;
    Kokkos::single(Kokkos::PerTeam(team),
                   [=]()
                   {
                     auto access = cs.access();
                     for (std::size_t j = 0; j < M::n_c; ++j)
                     {
                       access(j, 0) += scratch(j);
                     }
                   });
  }
 
  KOKKOS_INLINE_FUNCTION
  void
  operator()(const Tag3D _tag, const TeamMember& team) const
  {
    (void)_tag;
 
    const std::size_t N = m_particle_per_team;
    const std::size_t p0 = team.league_rank() * N;
 
    const auto _ntot = m_particles.n_particles();
    KOKKOS_ASSERT(p0 < _ntot)
    const auto upper_bound = ((p0 + N) >= _ntot) ? _ntot - p0 : N;
 
    KOKKOS_ASSERT(upper_bound > 0 && upper_bound <= _ntot);
 
    const std::size_t p = team.team_size();
 
    const std::size_t m = (upper_bound + p - 1) / p;
 
    const auto& positions = m_particles.position;
    const auto& contribs = m_particles.contribs;
    const auto& status = m_particles.status;
 
    Kokkos::parallel_for(
        Kokkos::TeamThreadRange(team, 0, upper_bound),
        [&](const std::size_t i)
        {
          const std::size_t flatten_index = p0 + i;
          KOKKOS_ASSERT(flatten_index < upper_bound);
          const bool active = status(flatten_index) == MC::Status::Idle;
          if (active)
          {
            const auto pos = m_particles.position(flatten_index);
            const auto weight = m_particles.get_weight(flatten_index);
            const auto new_status = M::update(random_pool,
                                              d_t,
                                              flatten_index,
                                              m_particles.model,
                                              contribs,
                                              pos,
                                              concentrations);
 
            if (new_status == MC::Status::Division)
            {
              division(flatten_index);
            }
            else
            {
              m_particles.ages(flatten_index, 1) += d_t;
            }
          }
        });
 
    team.team_barrier();
 
    Kokkos::parallel_for(
        Kokkos::TeamThreadRange(team, 0, p),
        [&](const std::size_t tid)
        {
          auto access = m_contribution_scatter.access();
          for (std::size_t k = 0; k < m; ++k)
          {
            const std::size_t idx = tid + k * p; // stride p
            const std::size_t flat_index = p0 + idx;
            if (idx >= upper_bound)
            {
              break;
            }
 
            KOKKOS_ASSERT(flat_index < _ntot);
 
            const double weight = m_particles.get_weight(p);
            const auto pos = positions(p);
            Kokkos::parallel_for(
                Kokkos::ThreadVectorRange(team, 0, M::n_c),
                [&](const int j)
                { access(j, pos) += weight * contribs(p, j); });
          }
        });
  }
};
 
#endif