move_kernel.hpp

#ifndef __SIMULATION_MOVE_KERNEL_HPP__
#define __SIMULATION_MOVE_KERNEL_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;
 
  template <class ExecutionSpace,
            class ViewType,
            class RandomPool,
            class IndexType = int64_t,
            const std::size_t CHUNK_SIZE>
  void
  fill_random(const ExecutionSpace& exec,
              ViewType a,
              RandomPool g,
              typename ViewType::const_value_type begin,
              typename ViewType::const_value_type end)
  {
    int64_t LDA = a.extent(0);
 
    if (LDA > 0)
    {
      Kokkos::parallel_for(
          "Kokkos::fill_random",
          Kokkos::RangePolicy<ExecutionSpace>(
              exec, 0, (LDA + (CHUNK_SIZE - 1)) / CHUNK_SIZE),
          Kokkos::Impl::fill_random_functor_begin_end<ViewType,
                                                      RandomPool,
                                                      CHUNK_SIZE,
                                                      ViewType::rank,
                                                      IndexType>(
              a, g, begin, end));
    }
  }

  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;
  }
  struct TagRNG
  {
  };
  struct TagMove
  {
  };
  struct TagLeave
  {
  };
 
  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),
          enable_move(b_move), 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;
      this->enable_leave = b_leave;
      this->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);
    }
 
    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, m),
                           [&rp, &rng, p, N](const std::size_t idx)
                           {
                             // Ok to use here, get_state should be called in
                             // each thread
                             auto gen = rp.get_state();
 
                             const std::size_t base = idx * p;
                             // current thread iteration p times with the same
                             // state
                             for (std::size_t k = 0; k < p; ++k)
                             {
                               const std::size_t i = base + k;
                               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;
                             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 TeamMember& team,
    //            std::size_t& local_dead_count) const
    // {
    //   (void)_tag;
    //   const std::size_t count = m_p_team_leave;
    //   const std::size_t p0 = team.league_rank() * count;
    //   const std::size_t n_particle = n_particles;
    //   const auto _d_t = static_cast<float>(d_t);
 
    //   const auto upper_bound
    //       = ((p0 + count) >= n_particle) ? n_particle - p0 : count;
    //   KOKKOS_ASSERT(upper_bound > 0 && upper_bound < n_particle);
 
    //   const std::size_t n_flow = move.leaving_flow.extent(0);
 
    //   using ScratchSpace = TeamPolicy::execution_space::scratch_memory_space;
    //   using ScratchView = Kokkos::View<MC::LeavingFlow*, ScratchSpace>;
    //   const auto leaving_flow = ScratchView(team.team_scratch(0), n_flow);
 
    //   // Kokkos::parallel_for(Kokkos::TeamVectorRange(team, n_flow),
    //   //                      [&](const std::size_t j)
    //   //                      { leaving_flow(j) = move.leaving_flow(j); });
 
    //   Kokkos::single(Kokkos::PerTeam(team),
    //                  [&]()
    //                  {
    //                    for (std::size_t j = 0; j < n_flow; ++j)
    //                    {
    //                      leaving_flow(j) = move.leaving_flow(j);
    //                    }
    //                  });
 
    //   team.team_barrier();
 
    //   std::size_t t_local = 0;
    //   Kokkos::parallel_reduce(
    //       Kokkos::TeamThreadRange(team, 0, upper_bound),
    //       [&](const std::size_t relative_index,
    //           std::size_t& thread_local_dead_count)
    //       {
    //         const std::size_t flatten_index = p0 + relative_index;
    //         if (status(flatten_index) == MC::Status::Idle)
    //         {
    //           ages(flatten_index, 0) += _d_t;
    //           handle_exit<ScratchSpace>(
    //               flatten_index, leaving_flow, thread_local_dead_count);
    //         }
    //       },
    //       t_local);
 
    //   team.team_barrier();
 
    //   Kokkos::single(Kokkos::PerTeam(team),
    //                  [&]()
    //                  {
    //                    Kokkos::single(Kokkos::PerThread(team),
    //                                   [&]() { local_dead_count += t_local;
    //                                   });
    //                  });
    //   // Kokkos::single(Kokkos::PerTeam(team),
    //   //                [&]()
    //   //                {
    //   //                  Kokkos::single(Kokkos::PerThread(team),
    //   //                                 [&]() {
    //   //                                 local_dead_count
    //   //                                 += t_local;
    //   //                                 });
    //   //                });
    // }
 
    KOKKOS_INLINE_FUNCTION void
    operator()([[maybe_unused]] TagLeave _tag,
               const std::size_t& idx,
               std::size_t& local_dead_count) const
    {
      if (status(idx) != MC::Status::Idle) [[unlikely]]
      {
        return;
      }
 
      ages(idx, 0) += d_t;
      handle_exit(idx, move.leaving_flow, local_dead_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 league_size = team_handle.league_size();
    //   const std::size_t league_rank = team_handle.league_rank();
    //   const std::size_t start_idx = league_rank * (n_particles /
    //   league_size); std::size_t end_idx = (league_rank + 1) * (n_particles /
    //   league_size);
 
    //   if (league_rank == (league_size - 1))
    //   {
    //     end_idx = n_particles;
    //   }
    //   Kokkos::parallel_for(
    //       Kokkos::TeamThreadRange(team_handle, start_idx, end_idx),
    //       [&](int idx)
    //       {
    //         ages(idx, 0) += d_t;
    //         handle_exit(idx, local_dead_count);
    //       });
    // }
 
    [[nodiscard]] bool
    need_launch() const
    {
      return enable_leave || enable_move;
    }
 
    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<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>();
        }
      }
    }
 
    // KOKKOS_INLINE_FUNCTION void
    // inner_handle_exit(const std::size_t i_flow,
    //                   const std::size_t idx,
    //                   const double liquid_volume,
    //                   const std::size_t position,
    //                   const MC::LeavingFlowView<true>& leaving_flow,
    //                   std ::size_t& dead_count) const
    // {
    //   // FIXME : when move AND exit, take the same random number,
    //   // is it really important ?
    //   const auto random_number = static_cast<float>(random(idx, 2));
    //   const auto& [index, flow] = leaving_flow(i_flow);
 
    //   const bool is_leaving = (position == index)
    //                           && probability_leaving<void>(
    //                               random_number, liquid_volume, flow, d_t);
 
    //   const int leave_mask = static_cast<int>(is_leaving);
 
    //   // If using probes
    //   if constexpr (AutoGenerated::FlagCompileTime::use_probe)
    //   {
    //     // Execute probe set, but only actually do something if leaving
    //     if (is_leaving)
    //     {
    //       const auto _ = probes.set(ages(idx, 0));
    //     }
    //   }
    //   if constexpr (AutoGenerated::FlagCompileTime::enable_event_counter)
    //   {
 
    //     events.add<MC::EventType::Exit>(leave_mask);
    //   }
 
    //   dead_count += leave_mask;
    //   ages(idx, 0) = (1 - leave_mask) * ages(idx, 0) /*leave_mask * 0 + */;
    //   status(idx) = is_leaving ? MC::Status::Exit : status(idx);
    // }
 
    // KOKKOS_FORCEINLINE_FUNCTION void
    // handle_exit(std::size_t idx,
    //             const MC::VolumeView<ComputeSpace, true>& liquid_volumes,
    //             std::size_t& dead_count) const
    // {
    //   const auto position = positions(idx);
    //   const auto liquid_volume = liquid_volumes(position);
    //   const MC::LeavingFlowView<true>& leaving_flow = move.leaving_flow;
    //   const std::size_t n_flow = leaving_flow.size();
    //   for (std::size_t i_flow = 0LU; i_flow < n_flow; ++i_flow)
    //   {
    //     inner_handle_exit(
    //         i_flow, idx, liquid_volume, position, leaving_flow, dead_count);
    //   }
    // }
 
    // KOKKOS_FORCEINLINE_FUNCTION void
    // handle_exit(std::size_t idx,
    //             const MC::VolumeView<ComputeSpace, true>& liquid_volumes,
    //             std::size_t& dead_count) const
    // {
    //   const std::size_t position = positions(idx);
    //   const double liquid_volume = liquid_volumes(position);
    //   const MC::LeavingFlowView<true>& leaving_flow = move.leaving_flow;
    //   const std::size_t n_flow = leaving_flow.size();
 
    //   // Strategy: most of the time there is only one flow  (0d reactor or
    //   // uniquement leaving point)
    //   //  Then the first flow is out of the loop
    //   //  For the first leaving flow use precomputed index_random_leave
    //   //  If the particle position is correct probability to leave is high
    //   //  then pregenerated only one number and get random on the fly in the
    //   //  loop if needed
 
    //   // One improvement is to use rng1 as long as as the we do consume it
    //   //  If first ok_p is false, rng1 is then not used
    //   //  This needs a branch ?
 
    //   auto rng1 = static_cast<float>(random(idx, index_random_leave));
 
    //   const auto& [index, flow] = leaving_flow(0);
    //   const bool p
    //       = probability_leaving<precision_tag>(rng1, liquid_volume, flow,
    //       d_t);
    //   const bool ok_p = position == index;
    //   const int m = static_cast<int>(ok_p) * static_cast<int>(p);
    //   int leave_mask = m;
 
    //   for (std::size_t i_flow = 1; i_flow < n_flow; ++i_flow)
    //   {
    //     if (leave_mask != 0)
    //     {
    //       break;
    //     }
    //     const auto& [index, flow] = leaving_flow(i_flow);
    //     const bool ok_p = position == index;
    //     if (!ok_p)
    //     {
    //       continue;
    //     }
 
    //     // if (!ok_p || leave_mask != 0)
    //     // {
    //     //   continue;
    //     // }
    //     auto gen = random_pool.get_state();
    //     rng1 = gen.frand(0, 1);
    //     random_pool.free_state(gen);
 
    //     const bool p = probability_leaving<precision_tag>(
    //         rng1, liquid_volume, flow, d_t);
 
    //     const int m = static_cast<int>(ok_p) * static_cast<int>(p);
    //     leave_mask |= m;
    //   }
 
    //   dead_count += leave_mask;
 
    //   // DO this betore age set to 0
    //   if constexpr (AutoGenerated::FlagCompileTime::use_probe)
    //   {
    //     // Execute probe set, but only actually do something if leaving
    //     if (leave_mask != 0)
    //     {
    //       // const auto _ = probes.template set<(ages(idx, 0));
    //     }
    //   }
 
    //   ages(idx, 0) *= (1 - leave_mask);
 
    //   // status(idx) = (leave_mask == 0) ? status(idx) : MC::Status::Exit;
 
    //   status(idx) = static_cast<MC::Status>(
    //       static_cast<int>(status(idx)) * (1 - leave_mask)
    //       + static_cast<int>(MC::Status::Exit) * leave_mask);
 
    //   if constexpr (AutoGenerated::FlagCompileTime::enable_event_counter)
    //   {
    //     events.add<MC::EventType::Exit>(leave_mask);
    //   }
    // }
    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
    {
 
      using mem_space = ComputeSpace::memory_space;
 
      const std::size_t position = positions(idx);
      // const MC::LeavingFlowView<true>& leaving_flow = move.leaving_flow;
      const std::size_t n_flow = leaving_flow.size();
      // const auto rng1 = static_cast<float>(random(idx, index_random_leave));
 
      // 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
      std::size_t i_flow = 0;
      double val_flow = 0.;
      double _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);
 
      int leave_mask = 0;
      // Cases
      // 0D: one flow and position always 0 then (val_flow != 0.) is always true
      // 3D: only for few particles
      //
      if (val_flow != 0.)
      {
        auto gen = random_pool.get_state();
        const auto rng1 = gen.frand(0., 1.);
        random_pool.free_state(gen);
 
        KOKKOS_ASSERT(_liquid_volume > 0.);
        KOKKOS_ASSERT(val_flow > 0.);
        const bool p = probability_leaving<precision_tag>(
            rng1, _liquid_volume, val_flow, d_t);
 
        leave_mask = static_cast<int>(p);
        // DO this betore age is reset to 0
        // if (p)
        // {
        dead_count += leave_mask;
        // }
        if constexpr (AutoGenerated::FlagCompileTime::use_probe)
        {
          if (leave_mask != 0)
          {
            const auto _ = probes.set<mem_space>(ages(idx, 0));
          }
        }
        ages(idx, 0) *= (1 - leave_mask);
        status(idx) = static_cast<MC::Status>(
            static_cast<int>(status(idx)) * (1 - leave_mask)
            + static_cast<int>(MC::Status::Exit) * leave_mask);
 
        if constexpr (AutoGenerated::FlagCompileTime::enable_event_counter)
        {
          events.add<MC::EventType::Exit>(leave_mask);
        }
      }
 
      return 0;
    }
 
    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{};
 
    // Kokkos::View<float**, Kokkos::LayoutLeft> random;
 
    bool enable_move{};
    bool enable_leave{};
  };
} // namespace Simulation::KernelInline
 
#endif