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<MoveLeaveTally, 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)
 
      // launch_move is skipped entirely when need_launch() is false, and
      // post_cycle tallies the reducer either way: clear it here so a step that
      // launches nothing cannot re-count the previous step's events.
      Kokkos::deep_copy(move_reducer, MoveLeaveTally{ 0, 0 });
 
      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 move_tally
          = Kokkos::create_mirror_view_and_copy(HostSpace(), move_reducer)();
 
      if constexpr (AutoGenerated::FlagCompileTime::enable_event_counter)
      {
        move_kernel.events.add<MC::EventType::Move>(move_tally.moved);
        move_kernel.events.add<MC::EventType::Exit>(move_tally.dead);
      }
 
      return std::tuple(host_red, move_tally.dead);
    }
 
    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,
                              _probes_div),
          m_options(options)
 
    {
    }

    void
    launch_move_3d(const std::size_t n_particle)
    {
      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 policy = Kokkos::TeamPolicy<TagMove>(move_space,
                                                static_cast<int>(league_size),
                                                Kokkos::AUTO(),
                                                Kokkos::AUTO());
 
      Kokkos::parallel_reduce("cycle_move", policy, move_kernel, move_reducer);
    }
 
    // Fused move + leave
    void
    launch_move_leave_3d(const std::size_t n_particle)
    {
      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 policy
          = Kokkos::TeamPolicy<TagMoveLeave>(move_space,
                                             static_cast<int>(league_size),
                                             Kokkos::AUTO(),
                                             Kokkos::AUTO());
 
      Kokkos::parallel_reduce(
          "cycle_move_leave", policy, move_kernel, move_reducer);
    }
 
    void
    launch_leave_0d(const std::size_t n_particle)
    {
      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);
 
      auto policy
          = Kokkos::TeamPolicy<TagLeaveB0D>(move_space,
                                            static_cast<int>(league_size),
                                            Kokkos::AUTO(),
                                            Kokkos::AUTO());
 
      Kokkos::parallel_reduce(
          "cycle_leave_0D", policy, move_kernel, move_reducer);
    }
 
    void
    launch_move(const std::size_t n_particle)
    {
      // The domain shape decides which kernels are legal:
      //   multi-compartment (3D) : move, or move + leave
      //   single compartment (0D): leave, or nothing
      // do_move() is exactly "the domain has more than one compartment".
      const bool is_multi_compartment = move_kernel.do_move();
      const bool has_outlet = move_kernel.do_leave();
 
      if (is_multi_compartment)
      {
        if (has_outlet)
        {
          launch_move_leave_3d(n_particle);
        }
        else
        {
          launch_move_3d(n_particle);
        }
        return;
      }
 
      if (has_outlet)
      {
        launch_leave_0d(n_particle);
      }
    }
 
    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()");
        KernelInline::CycleReducer<ComputeSpace> reducer(cycle_reducer);
        if (f_multi_compartment)
        {
 
          const auto policy_contribs
              = Kokkos::TeamPolicy<typename ContributionFunctor<Model>::Tag3D>(
                  model_space, league_size, Kokkos::AUTO(), Kokkos::AUTO());
          Kokkos::parallel_reduce("cycle_model_contribs",
                                  policy_contribs,
                                  contribution_kernel,
                                  reducer);
        }
        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_reduce("cycle_model_contribs_0d",
                                  policy_contribs,
                                  contribution_kernel,
                                  reducer);
        }
      }
    }
  };
 
} // namespace Simulation::KernelInline
 
#endif