model_kernel.hpp

#ifndef __SIMULATION_MC_KERNEL_HPP
#define __SIMULATION_MC_KERNEL_HPP
 
#include <Kokkos_Assert.hpp>
#include <Kokkos_Core.hpp>
#include <Kokkos_Macros.hpp>
#include <Kokkos_Printf.hpp>
#include <Kokkos_Random.hpp>
#include <biocma_cst_config.hpp>
#include <cassert>
#include <common/common.hpp>
#include <mc/alias.hpp>
#include <mc/domain.hpp>
#include <mc/events.hpp>
#include <mc/particles_container.hpp>
#include <mc/prng/prng.hpp>
#include <mc/traits.hpp>
// #include <simulation/probability_leaving.hpp>
#include <simulation/probe.hpp>
#include <utility>
 
#define CHECK_STATUS_OR_RETURN(__idx__)                                        \
  if (particles.status(__idx__) != MC::Status::Idle) [[unlikely]]              \
  {                                                                            \
    return;                                                                    \
  }
 
namespace Simulation::KernelInline
{
  struct TagContribution
  {
  };
  struct TagContribution0D
  {
  };
  struct TagCycle
  {
  };
 
  struct CycleReduceType
  {
    std::size_t waiting_allocation_particle;
    std::size_t dead_total;
 
    KOKKOS_INLINE_FUNCTION CycleReduceType&
    operator+=(const CycleReduceType& a)
    {
      this->waiting_allocation_particle += a.waiting_allocation_particle;
      this->dead_total += a.dead_total;
      return *this;
    }
  };
 
  template <class Space> class CycleReducer
  {
  public:
    // Required for Concept
    using reducer = CycleReducer;
    using value_type = CycleReduceType;
    using result_view_type = Kokkos::View<value_type, Space>;
 
    KOKKOS_INLINE_FUNCTION
    void
    join(value_type& dest, const value_type& src) const
    {
      // dest.dead_total += src.dead_total;
      // dest.waiting_allocation_particle += src.waiting_allocation_particle;
      dest += src;
    }
 
    [[nodiscard]] KOKKOS_INLINE_FUNCTION value_type&
    reference() const
    {
      return *value.data();
    }
 
    KOKKOS_INLINE_FUNCTION
    result_view_type
    view() const
    {
      return value;
    }
 
    [[nodiscard]] KOKKOS_INLINE_FUNCTION bool
    references_scalar() const
    {
      return references_scalar_v;
    }
 
    // Optional
    KOKKOS_INLINE_FUNCTION
    void
    init(value_type& val) const
    {
      val.dead_total = 0;
      val.waiting_allocation_particle = 0;
    }
 
    // KOKKOS_INLINE_FUNCTION
    // void final(value_type& val) const
    // {
    //   // NOP
    // }
 
    // Part of Build-In reducers for Kokkos
    KOKKOS_INLINE_FUNCTION
    explicit CycleReducer(value_type& value_)
        : value(&value_), references_scalar_v(true)
    {
    }
 
    KOKKOS_INLINE_FUNCTION
    explicit CycleReducer(const result_view_type& value_)
        : value(value_), references_scalar_v(false)
    {
    }
 
  private:
    result_view_type value;
    bool references_scalar_v;
  };
 
  template <ModelType M> struct CycleFunctor
  {
    using TeamPolicy = Kokkos::TeamPolicy<ComputeSpace>;
    using TeamMember = TeamPolicy::member_type;
    using value_type = CycleReduceType;
 
    CycleFunctor() = default;
 
    std::size_t n_p{};
 
    std::size_t m_p_team;
 
    KOKKOS_INLINE_FUNCTION
    CycleFunctor(std::size_t p_per_team,
                 MC::ParticlesContainer<M> _particles,
                 MC::pool_type _random_pool,
                 MC::KernelConcentrationType _concentrations,
                 MC::EventContainer _event,
                 ProbeAutogeneratedBuffer _probes)
        : m_p_team(p_per_team), d_t(0.), particles(std::move(_particles)),
          random_pool(std::move(_random_pool)),
          concentrations(std::move(_concentrations)), events(std::move(_event)),
          probes(std::move(_probes))
    {
    }
 
    [[nodiscard]] constexpr bool
    do_contribs() const
    {
      return M::n_c > 0;
    }
 
    void
    update(double _d_t, MC::ParticlesContainer<M> _particles)
    {
      this->d_t = _d_t;
      this->particles = std::move(_particles);
      n_p = this->particles.n_particles();
    }
 
    KOKKOS_INLINE_FUNCTION void
    operator()(TagCycle _tag,
               const TeamMember& team,
               value_type& reduce_val) const
    {
 
      (void)_tag;
      const std::size_t count = m_p_team;
      const std::size_t p0 = team.league_rank() * count;
      const std::size_t n_particle = n_p;
      const auto _d_t = d_t;
      const auto& status = particles.status;
 
      const auto upper_bound
          = ((p0 + count) >= n_particle) ? n_particle - p0 : count;
 
      KOKKOS_ASSERT(upper_bound > 0);
      KOKKOS_ASSERT(upper_bound <= n_particle);
 
      const auto& model = particles.model;
      const auto& contribs = particles.contribs;
 
      value_type local;
      Kokkos::parallel_reduce(
          Kokkos::TeamThreadRange(team, 0, upper_bound),
          [&](std::size_t relative_index, value_type& lv)
          {
            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 = particles.position(flatten_index);
              const auto new_status = M::update(random_pool,
                                                d_t,
                                                flatten_index,
                                                model,
                                                contribs,
                                                pos,
                                                concentrations);
 
              if (new_status == MC::Status::Division)
              {
                division(flatten_index, lv);
              }
              else
              {
                particles.ages(flatten_index, 1) += d_t;
              }
            }
          },
          local);
 
      team.team_barrier();
 
      Kokkos::single(Kokkos::PerTeam(team), [&]() { reduce_val += local; });
    }
 
    KOKKOS_INLINE_FUNCTION void
    division(const std::size_t idx, value_type& reduce_val) 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>(
            particles.ages(idx, 1)); // Skip error
      }
 
      if (!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>();
 
      /*
      TODO, it seems that after some calculation (example cstr 0d)
      size(division probes) = number of tally+1 where the last probe is 0
      where size(leaving time) == tally
      Is it a bug ?
      */
    }
 
    M::FloatType d_t;
    MC::ParticlesContainer<M> particles;
    MC::pool_type random_pool;
    MC::KernelConcentrationType concentrations;
    MC::EventContainer events;
    ProbeAutogeneratedBuffer probes;
  };
 
} // namespace Simulation::KernelInline
#endif