move_kernel.hpp

#ifndef __SIMULATION_MOVE_KERNEL_HPP__
#define __SIMULATION_MOVE_KERNEL_HPP__
 
#include "Kokkos_Macros.hpp"
#include "decl/Kokkos_Declare_OPENMP.hpp"
#include <Kokkos_Assert.hpp>
#include <Kokkos_Core.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/prng/prng.hpp>
#include <mc/traits.hpp>
#include <simulation/probability_leaving.hpp>
#include <simulation/probe.hpp>
#include <utility>
 
namespace Simulation::KernelInline
{
  constexpr bool enable_leave = true;
  constexpr bool disable_leave = false;
  constexpr bool disable_move = false;
  constexpr bool enable_move = true;
 
  struct TagMove
  {
  };
 
  struct TagLeave
  {
  };
 
  struct TagLeaveB0D
  {
  };

  KOKKOS_INLINE_FUNCTION std::size_t
  __find_next_compartment(
      const bool do_serch,
      const MC::NeighborsView<ComputeSpace, true>& neighbors,
      const MC::CumulativeProbabilityView<ComputeSpace, true>&
          cumulative_probability,
      const std::size_t i_compartment,
      const double random_number)
  {
    const int mask_do_serch = static_cast<int>(do_serch);
    const int max_neighbor = static_cast<int>(neighbors.extent(1));
 
    KOKKOS_ASSERT(max_neighbor >= 1);
    KOKKOS_ASSERT(random_number <= 1. && random_number >= 0.);
    KOKKOS_ASSERT(neighbors.extent(1) == cumulative_probability.extent(1));
 
    // Do not use cumulative_probability(i_compartment, max_neighbor - 1)==1
    // Bcause proba can be 0.999999999
    // KOKKOS_ASSERT(cumulative_probability(i_compartment, max_neighbor -
    // 1)>0.95); //5% relative
 
    // v1.1: This assert was there to be sure that the particle will leave but
    // actually
    // not needed because if rand< then left will be the last at the the ned of
    // the loop
    // KOKKOS_ASSERT(cumulative_probability(i_compartment, max_neighbor - 1)
    //               >= random_number);
 
    int left = 0;
    int right = mask_do_serch * (max_neighbor - 1);
    while (left < right)
    {
      const int mid = (left + right) >> 1; // NOLINT
      const auto pm = cumulative_probability(i_compartment, mid);
      const int mask = static_cast<int>(random_number > pm);
      left = mask * (mid + 1) + (1 - mask) * left;
      right = mask * right + (1 - mask) * mid;
    }
    KOKKOS_ASSERT(left >= 0 && static_cast<size_t>(left) < neighbors.extent(1));
    const auto ret = i_compartment * (1 - mask_do_serch)
                     + neighbors(i_compartment, left) * mask_do_serch;
    return ret;
  }
 
  template <typename ViewType1>
  KOKKOS_INLINE_FUNCTION void
  find_flow(const ViewType1& leaving_flow,
            const std::size_t position,
            MC::LeavingFlow::float_type& val_flow,
            MC::LeavingFlow::float_type& _liquid_volume)
  {
    std::size_t i_flow = 0;
    const std::size_t n_flow = leaving_flow.size();
    val_flow = 0.;
    _liquid_volume = 0.;
    // do-while because n_flow is likely to be 1
    do
    {
      const auto& [index, flow, liquid_volume] = leaving_flow(i_flow++);
      if (position == index)
      {
        val_flow = flow;
        _liquid_volume = liquid_volume;
        break;
      }
    } while (i_flow < n_flow);
  }
 
  struct MoveFunctor
  {
    using TeamPolicy = Kokkos::TeamPolicy<ComputeSpace>;
    using TeamMember = TeamPolicy::member_type;
    MoveFunctor() = default;
    MoveFunctor(std::size_t p_team_move,
                std::size_t p_team_leave,
                MC::ParticlePositions p,
                MC::ParticleStatus _status,
                MC::DomainState<ComputeSpace> m,
                MC::pool_type _random_pool,
                MC::EventContainer _events,
                ProbeAutogeneratedBuffer _probes,
                MC::ParticleAges _ages)
        : MoveFunctor(p_team_move,
                      p_team_leave,
                      0,
                      std::move(p),
                      std::move(_status),
                      0,
                      std::move(m),
                      std::move(_random_pool),
                      std::move(_events),
                      std::move(_probes),
                      std::move(_ages),
                      false,
                      false) {};
 
    MoveFunctor(std::size_t p_team_move,
                std::size_t p_team_leave,
                double _d_t,
                MC::ParticlePositions p,
                MC::ParticleStatus _status,
                std::size_t n_p,
                MC::DomainState<ComputeSpace> m,
                MC::pool_type _random_pool,
                MC::EventContainer _events,
                ProbeAutogeneratedBuffer _probes,
                MC::ParticleAges _ages,
                bool b_move,
                bool b_leave)
        : d_t(_d_t), positions(std::move(p)), n_particles(n_p),
          move(std::move(m)), random_pool(_random_pool), // NOLINT
          status(std::move(_status)), events(std::move(_events)),
          probes(std::move(_probes)), ages(std::move(_ages)),
          m_p_team_leave(p_team_leave), m_p_team_move(p_team_move),
          m_enable_move(b_move), m_enable_leave(b_leave) {};
 
    void
    update(double _d_t,
           std::size_t n_p,
           MC::DomainState<ComputeSpace> move_i,
           MC::ParticlePositions _positions,
           MC::ParticleStatus _status,
           MC::ParticleAges _ages,
           bool b_move,
           bool b_leave)
    {
 
      this->d_t = _d_t;
      this->n_particles = n_p;
      m_enable_leave = b_leave;
      m_enable_move = b_move;
      this->move = std::move(move_i);
      this->positions = std::move(_positions);
      this->status = std::move(_status);
      this->ages = std::move(_ages);
    }
 
    [[nodiscard]] bool
    need_launch() const
    {
      return m_enable_leave || m_enable_move;
    }
 
    KOKKOS_INLINE_FUNCTION void
    operator()(TagMove /*tag*/,
               const Kokkos::TeamPolicy<ComputeSpace>::member_type& team) const
    {
      using ScratchSpace
          = Kokkos::TeamPolicy<>::execution_space::scratch_memory_space;
      using ScratchView = Kokkos::View<float*, ScratchSpace>;
 
      const std::size_t count = m_p_team_move;
      const std::size_t p0 = team.league_rank() * count;
      const std::size_t n_particle = n_particles;
 
      const auto upper_bound
          = ((p0 + count) >= n_particle) ? n_particle - p0 : count;
      KOKKOS_ASSERT(upper_bound > 0 && upper_bound < n_particle);
      const auto& rp = random_pool;
 
      const std::size_t N = count * 2;
 
      // rng is actually a flat m*p array
      const std::size_t p = team.team_size();
      const std::size_t m = (N + p - 1) / p;
 
      ScratchView rng(team.team_scratch(0), N);
 
      // Use "tiling" to minimize contention when aquired_state
      // State is aquired m times instead of N, it is supposed to reduce
      // contention
 
      Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 0, p),
                           [&rp, &rng, N, m, p](const std::size_t idx)
                           {
                             // Ok to use here, get_state should be called in
                             // each thread
                             auto gen = rp.get_state();
 
                             // current thread iteration p times with the same
                             // state
                             for (std::size_t k = 0; k < m; ++k)
                             {
                               const std::size_t i = idx + k * p; // stride p
 
                               if (i >= N)
                               {
                                 break;
                               }
 
                               rng(i) = gen.frand(0., 1.);
                             }
 
                             rp.free_state(gen);
                           });
      team.team_barrier();
 
      // We can use flat array index here ordering of random doesnt matter
      Kokkos::parallel_for(Kokkos::TeamThreadRange(team, 0, upper_bound),
                           [&](const std::size_t idx)
                           {
                             const auto flat_index = p0 + idx;
 
                             if (status(flat_index) == MC::Status::Idle)
                             {
 
                               const std::size_t base = idx * 2;
                               KOKKOS_ASSERT(base + 1 < N);
                               const auto rng1 = rng(base);
                               const auto rng2 = rng(base + 1);
                               handle_move(flat_index, rng1, rng2);
                             }
                           });
    }
 
    KOKKOS_INLINE_FUNCTION void
    operator()(TagLeave _tag,
               const Kokkos::TeamPolicy<ComputeSpace>::member_type& team_handle,
               std::size_t& local_dead_count) const
    {
      (void)_tag;
      const std::size_t count = m_p_team_leave;
      const std::size_t p0 = team_handle.league_rank() * count;
      const std::size_t n_particle = n_particles;
 
      const auto upper_bound
          = ((p0 + count) >= n_particle) ? n_particle - p0 : count;
 
      KOKKOS_ASSERT(upper_bound > 0 && upper_bound <= n_particle);
      std::size_t team_dead_count = 0;
      const auto& lf = move.leaving_flow;
      Kokkos::parallel_reduce(
          Kokkos::TeamThreadRange(team_handle, 0, upper_bound),
          [&](const std::size_t idx, std::size_t& thread_dead_count)
          {
            const auto flat_index = p0 + idx;
            // ages(flat_index, 0) += d_t;
            handle_exit(flat_index, lf, thread_dead_count);
          },
          team_dead_count);
 
      team_handle.team_barrier();
      Kokkos::single(
          Kokkos::PerTeam(team_handle),
          [&]()
          {
            if constexpr (AutoGenerated::FlagCompileTime::enable_event_counter)
            {
              events.add<MC::EventType::Exit>(team_dead_count);
            }
            local_dead_count += team_dead_count;
          });
    }
 
    bool
    do_move() const
    {
      return this->m_enable_move;
    }
 
    bool
    do_leave() const
    {
      return m_enable_leave;
    }
 
    KOKKOS_INLINE_FUNCTION void
    operator()(TagLeaveB0D _tag,
               const Kokkos::TeamPolicy<ComputeSpace>::member_type& team_handle,
               std::size_t& local_dead_count) const
    {
 
      (void)_tag;
 
      const auto& lf = move.leaving_flow;
 
      // In 0d, the flow is the same for every particle
      // Only one member of the team calculate lambda and
      // broadcast it to other members
      double lambda = 0.;
      Kokkos::single(
          Kokkos::PerTeam(team_handle),
          [&lf, dt = this->d_t](double& local_lambda)
          {
            const auto& [_, flow_value, liquid_volume] = lf(0);
 
            local_lambda = dt * flow_value / liquid_volume;
          },
          lambda);
 
      team_handle.team_barrier();
      // In 0d context we can even be sure that flow will never be 0
      // Note: if flow is 0, proba will be 0 and then the result will be the
      // same Can we remove the condition?
      //
      if (lambda == 0.)
      {
        return;
      }
      const std::size_t N = m_p_team_leave;
      const std::size_t p0 = team_handle.league_rank() * N;
      const auto upper_bound = ((p0 + N) >= n_particles) ? n_particles - p0 : N;
 
      KOKKOS_ASSERT(upper_bound > 0 && upper_bound <= n_particles);
 
      const auto& rp = random_pool;
 
      const std::size_t p = team_handle.team_size();
 
      const std::size_t m = (upper_bound + p - 1) / p;
 
      std::size_t team_dead_count = 0;
 
      // For each thread of the team we assigm m iterations
      // This divide the number of call of get_state by m and may reduce
      // contention especially on CPU typically P = 1 on cpu then m is
      // equal to npt (128-4096)
      Kokkos::parallel_reduce(
          Kokkos::TeamThreadRange(team_handle, 0, p),
          [&](const std::size_t tid, std::size_t& dead_count)
          {
            auto gen = rp.get_state();
 
            for (std::size_t k = 0; k < m; ++k)
            {
              const std::size_t idx = tid + k * p; // stride p
 
              if (idx >= upper_bound)
              {
                break;
              }
              const std::size_t flat_index = p0 + idx;
 
              if (status(flat_index) != MC::Status::Idle)
              {
                continue;
              }
              const double r = gen.drand(0., 1.);
 
              perform_exit(
                  probability_leaving<decltype(r), precision_tag>(r, lambda),
                  flat_index,
                  dead_count);
            }
            rp.free_state(gen);
          },
          team_dead_count);
 
      team_handle.team_barrier();
 
      Kokkos::single(
          Kokkos::PerTeam(team_handle),
          [&]()
          {
            if constexpr (AutoGenerated::FlagCompileTime::enable_event_counter)
            {
              events.add<MC::EventType::Exit>(team_dead_count);
            }
            local_dead_count += team_dead_count;
          });
    }
 
    KOKKOS_FUNCTION void
    handle_move(const std::size_t idx, const float rng1, const float rng2) const
    {
 
      // const auto rng1 = static_cast<float>(random(idx, 0));
      // const auto rng2 = static_cast<float>(random(idx, 1));
 
      KOKKOS_ASSERT(rng1 >= 0. && rng1 <= 1 && rng2 >= 0. && rng2 <= 1);
 
      const std::size_t i_current_compartment = positions(idx);
 
      KOKKOS_ASSERT(
          i_current_compartment < move.liquid_volume.extent(0)
          && "Particle position is incorect (greater than compartment number)");
 
      const bool mask_next = probability_leaving<float, fast_tag>(
          rng1,
          move.liquid_volume(i_current_compartment),
          move.diag_transition(i_current_compartment),
          d_t);
 
      positions(idx) = __find_next_compartment(mask_next,
                                               move.neighbors,
                                               move.cumulative_probability,
                                               i_current_compartment,
                                               rng2);
 
      // positions(idx)
      //     = (mask_next) ? __find_next_compartment(move.neighbors,
      //                                             move.cumulative_probability,
      //                                             i_current_compartment,
      //                                             rng2)
      //                   : i_current_compartment;
 
      KOKKOS_ASSERT(
          positions(idx) < move.liquid_volume.extent(0)
          && " Position after move is greater than compartment number");
 
      if constexpr (AutoGenerated::FlagCompileTime::enable_event_counter)
      {
        if (mask_next)
        {
          events.wrap_incr<MC::EventType::Move>();
        }
      }
    }
 
    template <typename ExecSpace>
    KOKKOS_FORCEINLINE_FUNCTION std::size_t
    handle_exit(
        const std::size_t idx,
        const Kokkos::View<const MC::LeavingFlow*, ExecSpace>& leaving_flow,
        std::size_t& dead_count) const
    {
 
      const std::size_t position = positions(idx);
      // Strategy:
      //  first find the value of leaving flow (0-> particle doesn´t leave)
      //  do-while +early break is ok as n_flow is likely <10
      // second: calculate probability leaving, flow=0 => p=0 theres no need
      // to check condition
      MC::LeavingFlow::float_type found_flow_value = 0.;
      MC::LeavingFlow::float_type found_liquid_volume = 0.;
      find_flow(leaving_flow, position, found_flow_value, found_liquid_volume);
 
      // Cases
      // only for few particles
      if (found_flow_value != 0.)
      {
 
        auto gen = random_pool.get_state();
        const auto r = gen.frand(0., 1.);
        random_pool.free_state(gen);
 
        KOKKOS_ASSERT(found_liquid_volume > 0.);
        KOKKOS_ASSERT(found_flow_value > 0.);
        const bool p = probability_leaving<decltype(r), precision_tag>(
            r, found_liquid_volume, found_flow_value, d_t);
        perform_exit(p, idx, dead_count);
      }
 
      return 0;
    }
 
    KOKKOS_INLINE_FUNCTION void
    perform_exit(const bool proba,
                 const std::size_t idx,
                 std::size_t& dead_count) const
    {
 
      if (proba)
      {
        ++dead_count;
 
        if constexpr (AutoGenerated::FlagCompileTime::use_probe)
        {
          using mem_space = ComputeSpace::memory_space;
          // DO this betore age is reset to 0
          const auto _ = probes.set<mem_space>(ages(idx, 0));
        }
 
        ages(idx, 0) = 0;
        status(idx) = MC::Status::Exit;
      }
      else
      {
        ages(idx, 0) += d_t;
      }
    }
 
    double d_t{};
    MC::ParticlePositions positions;
    std::size_t n_particles{};
    MC::DomainState<ComputeSpace, true> move;
    MC::pool_type random_pool;
    MC::ParticleStatus status;
    MC::EventContainer events;
    ProbeAutogeneratedBuffer probes;
    MC::ParticleAges ages;
    std::size_t m_p_team_leave{};
    std::size_t m_p_team_move{};
 
    bool m_enable_move{};
    bool m_enable_leave{};
  };
} // namespace Simulation::KernelInline
 
#endif