kernels.hpp

#ifndef __SIMULATION_KERNELS_HPP__
#define __SIMULATION_KERNELS_HPP__
 
#include <Kokkos_Core_fwd.hpp>
#include <common/common.hpp>
#include <common/kokkos_getpolicy.hpp>
#include <mc/domain.hpp>
#include <mc/unit.hpp>
#include <simulation/kernels/contribution_kernel.hpp>
#include <simulation/kernels/model_kernel.hpp>
#include <simulation/kernels/move_kernel.hpp>
 
#include <common/execinfo.hpp>
#include <stdexcept>
 
namespace Simulation::KernelInline
{
 
  template <typename Space>
  using cycle_reducer_view_type
      = KernelInline::CycleReducer<Space>::result_view_type;
  template <typename Space>
  using move_reducer_view_type = Kokkos::View<std::size_t, Space>;
 
  template <typename Space, ModelType Model> struct CycleFunctors
  {
    using FModel = Model;
 
    using cycle_kernel_type = CycleFunctor<Model>;
    using move_kernel_type = MoveFunctor;
 
    ComputeSpace move_space;
 
    ComputeSpace model_space;
 
    cycle_reducer_view_type<Space> cycle_reducer;
    move_reducer_view_type<Space> move_reducer;
    cycle_kernel_type cycle_kernel;
    move_kernel_type move_kernel;
 
    ContributionFunctor<Model> contribution_kernel;
 
    CycleFunctors() = default;
 
    KernelDispatchOptions m_options{};
 
    bool f_multi_compartment;
 
    void
    update(const double d_t,
           MC::ParticlesContainer<Model> container,
           MC::DomainState<ComputeSpace>&& new_move)
    {
      f_multi_compartment = new_move.liquid_volume.size() > 1;
      const bool enable_move = f_multi_compartment;
      const bool enable_leave = new_move.leaving_flow.size() != 0;
 
      // FIXME: cycle_kernel: need to update container because we change:
      // n_used_element counter which is size_t outside of kernel. As functor
      // owns a copy of container this counter is not update between iterations.
      // 1. Use n_used_element as a reference counter type
      // 2. Manually update counter in update function (dirty way)
 
      cycle_kernel.update(d_t, container);
 
      contribution_kernel.update(container);
 
      // TODO: Why need to update all views (where did we lost the refcount ? )
      move_kernel.update(d_t,
                         container.n_particles(),
                         std::move(new_move),
                         container.position,
                         container.status,
                         container.ages,
                         enable_move,
                         enable_leave);
    }
 
    auto
    get_host_reduction()
    {
      const auto host_red
          = Kokkos::create_mirror_view_and_copy(HostSpace(), cycle_reducer)();
 
      const auto host_out_counter
          = Kokkos::create_mirror_view_and_copy(HostSpace(), move_reducer)();
 
      return std::tuple(host_red, host_out_counter);
    }
 
    CycleFunctors(KernelDispatchOptions options,
                  MC::ParticlesContainer<Model> container,
                  MC::pool_type _random_pool,
                  MC::KernelConcentrationType _concentrations,
                  MC::ContributionView _contribs_scatter,
                  MC::EventContainer _event,
                  MC::DomainState<ComputeSpace> m,
                  ProbeAutogeneratedBuffer _probes,
                  ProbeAutogeneratedBuffer _probes_div)
        : cycle_reducer("cycle_reducer"), move_reducer("move_reducer"),
          cycle_kernel(options.m_p_p_team_model,
                       container,
                       _random_pool,
                       std::move(_concentrations),
                       _event,
                       _probes_div),
          move_kernel(options.m_p_p_team_move,
                      options.m_p_p_team_leave,
                      container.position,
                      container.status,
                      std::move(m),
                      _random_pool,
                      _event,
                      _probes,
                      container.ages),
          contribution_kernel(
              options.m_p_p_team_contribs, _contribs_scatter, container),
          m_options(options)
 
    {
    }
 
    void
    launch_move(const std::size_t n_particle) const
    {
 
      if (move_kernel.enable_move)
      {
        const auto npt = m_options.m_p_p_team_move;
        if (n_particle <= npt)
        {
          // TODO
          throw std::runtime_error("Nparticle<n per team");
        }
 
        const std::size_t league_size = Common::c_league_size(n_particle, npt);
 
        auto cycle_policy
            = Kokkos::TeamPolicy<TagMove>(model_space,
                                          static_cast<int>(league_size),
                                          Kokkos::AUTO(),
                                          Kokkos::AUTO());
 
        cycle_policy.set_scratch_size(0,
                                      Kokkos::PerTeam(sizeof(float) * npt * 2));
 
        Kokkos ::parallel_for("cycle_move", cycle_policy, move_kernel);
      }
 
      if (move_kernel.enable_leave)
      {
 
        const auto _policy_leave = Kokkos::RangePolicy<KernelInline::TagLeave>(
            move_space, 0, n_particle);
        Kokkos ::parallel_reduce(
            "cycle_move_leave", _policy_leave, move_kernel, move_reducer);
      }
    }
 
    void
    launch_model(const std::size_t n_particle) const
    {
      if (n_particle <= m_options.m_p_p_team_model)
      {
        // TODO
        throw std::runtime_error("Nparticle<n per team");
      }
 
      std::size_t league_size
          = Common::c_league_size(n_particle, m_options.m_p_p_team_model);
 
      const auto cycle_policy
          = Kokkos::TeamPolicy<TagCycle>(model_space,
                                         static_cast<int>(league_size),
                                         Kokkos::AUTO(),
                                         Kokkos::AUTO());
 
      Kokkos::parallel_reduce(
          "cycle_model",
          cycle_policy,
          cycle_kernel,
          KernelInline::CycleReducer<ComputeSpace>(cycle_reducer));
      Kokkos::fence(); // TODO needed ?
 
      // Assumptions
      // Newborn cells do not contribte in current time step
      // Mother cell doesn´t exist but
      // Contribution array is not changed during division and
 
      if (cycle_kernel.do_contribs())
      {
 
        if (n_particle <= m_options.m_p_p_team_contribs)
        {
          // TODO
          throw std::runtime_error("Nparticle<n per team");
        }
 
        league_size
            = Common::c_league_size(n_particle, m_options.m_p_p_team_contribs);
        static_assert(ConstWeightModelType<Model>,
                      "ModelType:Constapply_weight()");
 
        if (f_multi_compartment)
        {
 
          const auto policy_contribs
              = Kokkos::TeamPolicy<typename ContributionFunctor<Model>::Tag3D>(
                  model_space, league_size, Kokkos::AUTO(), Kokkos::AUTO());
          Kokkos::parallel_for(
              "cycle_model_contribs", policy_contribs, contribution_kernel);
        }
        else
        {
          auto policy_contribs
              = Kokkos::TeamPolicy<typename ContributionFunctor<Model>::Tag0D>(
                  model_space, league_size, Kokkos::AUTO(), Kokkos::AUTO());
          policy_contribs.set_scratch_size(
              0, Kokkos::PerTeam(sizeof(float) * Model::n_c));
          Kokkos::parallel_for(
              "cycle_model_contribs_0d", policy_contribs, contribution_kernel);
        }
      }
    }
  };
 
} // namespace Simulation::KernelInline
 
#endif