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/unit.hpp>
#include <simulation/kernels/model_kernel.hpp>
#include <simulation/kernels/move_kernel.hpp>
 
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;
 
    CycleFunctors() = default;
 
    void
    update(const double d_t,
           MC::ParticlesContainer<Model> container,
           MC::DomainState<ComputeSpace>&& new_move)
    {
      const bool enable_move = new_move.liquid_volume.size() > 1;
      const bool enable_leave = new_move.leaving_flow.size() != 0;
 
      auto n_sample = enable_move ? 2 : 0;
      n_sample += enable_leave ? 1 : 0;
      container.change_nsample(n_sample);
 
      // 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);
 
      // TODO: Why need to update all views (where did we lost the refcount ? )
      move_kernel.update(move_space,
                         d_t,
                         container.n_particles(),
                         std::move(new_move),
                         container.position,
                         container.status,
                         container.ages,
                         container.random,
                         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(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(container,
                       _random_pool,
                       std::move(_concentrations),
                       _contribs_scatter,
                       _event,
                       _probes_div),
          move_kernel(container.position,
                      container.status,
                      std::move(m),
                      _random_pool,
                      _event,
                      _probes,
                      container.ages,
                      container.random)
    {
    }
 
    void
    launch_move(const std::size_t n_particle) const
    {
 
      if (move_kernel.enable_move)
      {
        // const auto _policy_move =
        //     Common::get_policy<move_kernel_type, KernelInline::TagMove>(
        //         move_kernel, n_particle, false);
        // const auto _policy_move =
        //     Common::get_policy<KernelInline::TagMove>(n_particle, false);
 
        const auto _policy_move = Kokkos::RangePolicy<KernelInline::TagMove>(
            move_space, 0, n_particle);
 
        Kokkos ::parallel_for("cycle_move", _policy_move, move_kernel);
      }
      if (move_kernel.enable_leave)
      {
 
        // const auto _policy_leave =
        //     Kokkos::TeamPolicy<ComputeSpace, KernelInline::TagLeave>(
        //         256, Kokkos::AUTO);
 
        // Kokkos ::parallel_reduce(
        //     "cycle_move_leave", _policy_leave, move_kernel, move_reducer);
 
        const auto _policy_leave = Kokkos::RangePolicy<KernelInline::TagLeave>(
            move_space, 0, n_particle);
        // const auto _policy_move = Kokkos::TeamPolicy<TagLeave>(56, 256);
        Kokkos ::parallel_reduce(
            "cycle_move_leave", _policy_leave, move_kernel, move_reducer);
      }
    }
 
    void
    launch_model(const std::size_t n_particle) const
    {
 
      auto _policy = Kokkos::RangePolicy<TagCycle>(model_space, 0, n_particle);
      Kokkos::parallel_reduce(
          "cycle_model",
          _policy,
          cycle_kernel,
          KernelInline::CycleReducer<ComputeSpace>(cycle_reducer));
      Kokkos::fence(); // TODO needed ?
 
      constexpr int PARTICLES_PER_TEAM = 256;
      int league_size = Kokkos::ceil(n_particle / PARTICLES_PER_TEAM);
 
      auto _policy2 = Kokkos::TeamPolicy<TagContribution>(
          model_space, league_size, Kokkos::AUTO(), Kokkos::AUTO());
      Kokkos::parallel_for("cycle_model2", _policy2, cycle_kernel);
    }
  };
 
} // namespace Simulation::KernelInline
 
#endif