kernels.hpp

#ifndef __SIMULATION_KERNELS_HPP__
#define __SIMULATION_KERNELS_HPP__
 
#include "mc/alias.hpp"
#include <Kokkos_Core_fwd.hpp>
#include <common/common.hpp>
#include <common/kokkos_getpolicy.hpp>
#include <impl/Kokkos_Profiling.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, d_t);
 
      // 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,
                       _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),
          contribution_kernel(options.m_p_p_team_contribs,
                              _contribs_scatter,
                              container,
                              _random_pool,
                              _concentrations,
                              _event),
          m_options(options)
 
    {
    }
 
    void
    launch_move(const std::size_t n_particle)
    {
      bool is_0d_reactor = move_kernel.do_move();
      if (is_0d_reactor)
      {
 
        std::size_t npt = m_options.m_p_p_team_move;
 
        if (n_particle <= npt)
        {
          move_kernel.m_p_team_move = 1;
          npt = 1;
        }
 
        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.do_leave())
      {
 
        std::size_t npt = move_kernel.m_p_team_leave;
        if (n_particle <= npt)
        {
          move_kernel.m_p_team_leave = 1;
          npt = 1;
        }
        const std::size_t league_size = Common::c_league_size(n_particle, npt);
 
        // This is duplicated code but there is no simple alternative to select
        // policy tag at runtime
        if (!is_0d_reactor)
        {
          auto _policy_leave
              = Kokkos::TeamPolicy<TagLeaveB0D>(move_space,
                                                static_cast<int>(league_size),
                                                Kokkos::AUTO(),
                                                Kokkos::AUTO());
 
          Kokkos::parallel_reduce(
              "cycle_move_leave", _policy_leave, move_kernel, move_reducer);
        }
        else
        {
          auto _policy_leave
              = Kokkos::TeamPolicy<TagLeave>(move_space,
                                             static_cast<int>(league_size),
                                             Kokkos::AUTO(),
                                             Kokkos::AUTO());
 
          // auto _policy_leave
          //     = Kokkos::RangePolicy<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)
    {
 
      std::size_t npt = m_options.m_p_p_team_model;
 
      if (n_particle <= npt)
      {
        cycle_kernel.m_p_team = 1;
        npt = 1;
      }
 
      std::size_t league_size = Common::c_league_size(n_particle, npt);
 
      // const auto cycle_policy
      //     = Kokkos::TeamPolicy<TagCycle, Kokkos::Schedule<Kokkos::Dynamic>>(
      //         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())
      {
 
        std::size_t npt = m_options.m_p_p_team_contribs;
        if (n_particle <= npt)
        {
          contribution_kernel.m_particle_per_team = 1;
          npt = 1;
        }
 
        league_size = Common::c_league_size(n_particle, npt);
        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