move_kernel.hpp

#ifndef __SIMULATION_MOVE_KERNEL_HPP__
#define __SIMULATION_MOVE_KERNEL_HPP__
 
#include "Kokkos_Macros.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
  {
  };

  struct TagMoveLeave
  {
  };

  struct MoveLeaveTally
  {
    std::size_t moved;
    std::size_t dead;
 
    KOKKOS_INLINE_FUNCTION MoveLeaveTally&
    operator+=(const MoveLeaveTally& other)
    {
      moved += other.moved;
      dead += other.dead;
      return *this;
    }
  };

  KOKKOS_INLINE_FUNCTION std::size_t
  __find_next_compartment(
      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 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 = 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));
    return neighbors(i_compartment, left);
  }
 
  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,
               MoveLeaveTally& local_tally) const
    {
      const std::size_t N = m_p_team_move;
      const std::size_t p0 = team.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 auto& proba = move.move_probability;
      const std::size_t p = team.team_size();
      const std::size_t m = (upper_bound + p - 1) / p;
 
      std::size_t team_move_count = 0;
      Kokkos::parallel_reduce(
          Kokkos::TeamThreadRange(team, 0, p),
          [&](const std::size_t tid, std::size_t& moved_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;
              }
 
              bool moved = false;
              const std::size_t next
                  = next_compartment(positions(flat_index), proba, gen, moved);
              if (moved)
              {
                positions(flat_index) = next;
                ++moved_count;
              }
            }
 
            rp.free_state(gen);
          },
          team_move_count);
 
      team.team_barrier();
      Kokkos::single(Kokkos::PerTeam(team),
                     [&]() { local_tally.moved += team_move_count; });
    }
 
    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;
          });
    }


    KOKKOS_INLINE_FUNCTION void
    operator()(TagMoveLeave /*tag*/,
               const Kokkos::TeamPolicy<ComputeSpace>::member_type& team,
               MoveLeaveTally& local_tally) const
    {
      const std::size_t N = m_p_team_move;
      const std::size_t p0 = team.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 auto& lf = move.leaving_flow;
      const auto& proba = move.move_probability;
 
      const std::size_t p = team.team_size();
      const std::size_t m = (upper_bound + p - 1) / p;
 
      MoveLeaveTally team_tally{ 0, 0 };
 
      Kokkos::parallel_reduce(
          Kokkos::TeamThreadRange(team, 0, p),
          [&](const std::size_t tid, MoveLeaveTally& acc)
          {
            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;
 
              const bool is_idle = status(flat_index) == MC::Status::Idle;
              std::size_t position = positions(flat_index);
 
              if (is_idle)
              {
                bool moved = false;
                const std::size_t next
                    = next_compartment(position, proba, gen, moved);
                // Only movers write back. Most particles stay put in a
                // step, so an unconditional store dirties every cache line of
                // `positions` to rewrite the value it already held.
                if (moved)
                {
                  positions(flat_index) = next;
                  position = next;
                  ++acc.moved;
                }
              }
 
              ages(flat_index, 0) += d_t;
 
              // Only idle particles may exit: an already-Exit one in an
              // outlet compartment would otherwise be recounted every step
              // and overshoot inactive_counter.
              if (is_idle)
              {
                handle_exit_at(flat_index, lf, position, gen, acc.dead);
              }
            }
 
            rp.free_state(gen);
          },
          team_tally);
 
      team.team_barrier();
 
      Kokkos::single(Kokkos::PerTeam(team),
                     [&]() { local_tally += team_tally; });
    }
 
    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,
               MoveLeaveTally& local_tally) 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.);
              ages(flat_index, 0) += d_t;
              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),
                     [&]() { local_tally.dead += team_dead_count; });
    }

    template <typename GenType>
    KOKKOS_FORCEINLINE_FUNCTION std::size_t
    next_compartment(
        const std::size_t i_current_compartment,
        const MC::MoveProbabilityView<ComputeSpace, true>& move_probability,
        GenType& gen,
        bool& moved) const
    {
      KOKKOS_ASSERT(
          i_current_compartment < move.liquid_volume.extent(0)
          && "Particle position is incorect (greater than compartment number)");
 
      moved = move_probability(i_current_compartment) > gen.frand(0.F, 1.F);
 
      if (!moved)
      {
        return i_current_compartment;
      }
 
      const auto next = __find_next_compartment(move.neighbors,
                                                move.cumulative_probability,
                                                i_current_compartment,
                                                gen.frand(0.F, 1.F));
      KOKKOS_ASSERT(
          next < move.liquid_volume.extent(0)
          && " Position after move is greater than compartment number");
      return next;
    }

    template <typename ExecSpace, typename GenType>
    KOKKOS_FORCEINLINE_FUNCTION void
    handle_exit_at(
        const std::size_t idx,
        const Kokkos::View<const MC::LeavingFlow*, ExecSpace>& leaving_flow,
        const std::size_t position,
        GenType& gen,
        std::size_t& dead_count) const
    {
      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);
 
      // Only a few particles sit in an outlet compartment
      if (found_flow_value != 0.)
      {
        const auto r = gen.frand(0., 1.);
        KOKKOS_ASSERT(found_liquid_volume > 0.);
        KOKKOS_ASSERT(found_flow_value > 0.);
        perform_exit(probability_leaving<decltype(r), precision_tag>(
                         r, found_liquid_volume, found_flow_value, d_t),
                     idx,
                     dead_count);
      }
    }
 
    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
 
namespace Kokkos
{
  template <>
  struct reduction_identity<Simulation::KernelInline::MoveLeaveTally>
  {
    KOKKOS_FORCEINLINE_FUNCTION static Simulation::KernelInline::MoveLeaveTally
    sum()
    {
      return { 0, 0 };
    }
  };
} // namespace Kokkos
 
#endif