move_kernel.hpp

#ifndef __SIMULATION_MOVE_KERNEL_HPP__
#define __SIMULATION_MOVE_KERNEL_HPP__
 
#include "Kokkos_Macros.hpp"
#include "common/common.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 <decl/Kokkos_Declare_OPENMP.hpp>
#include <mc/alias.hpp>
#include <mc/domain.hpp>
#include <mc/events.hpp>
#include <mc/prng/prng.hpp>
#include <mc/traits.hpp>
// #include <simulation/move_info.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;
 
  // KOKKOS_INLINE_FUNCTION std::size_t __find_next_compartment(
  //     const ConstNeighborsView<ComputeSpace>& neighbors,
  //     const CumulativeProbabilityView<ComputeSpace>& cumulative_probability,
  //     const std::size_t i_compartment,
  //     const double random_number)
  // {
  //   const int max_neighbor = static_cast<int>(neighbors.extent(1));
 
  //   std::size_t next =
  //       neighbors(i_compartment, 0); // Default to the first neighbor
 
  //   // Iterate through the neighbors to find the appropriate next compartment
  //   for (int k_neighbor = 0; k_neighbor < max_neighbor - 1; ++k_neighbor)
  //   {
 
  //     // Get the cumulative probability range for the current neighbor
  //     const auto pi = cumulative_probability(i_compartment, k_neighbor);
  //     const auto pn = cumulative_probability(i_compartment, k_neighbor + 1);
 
  //     // Use of a Condition mask to avoid branching.
  //     next = (random_number <= pn && pi <= random_number)
  //                ? neighbors(i_compartment, k_neighbor + 1)
  //                : next;
  //   }
 
  //   return next; // Return the index of the chosen next compartment
  // }

  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));
 
    int left = 0;
    int right = mask_do_serch * (max_neighbor - 1);
    while (left < right)
    {
      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;
    }
 
    const auto ret = i_compartment * (1 - mask_do_serch)
                     + neighbors(i_compartment, left) * mask_do_serch;
    return ret;
  }
 
  // KOKKOS_INLINE_FUNCTION std::size_t __find_next_compartment(
  //     const bool do_serch,
  //     const ConstNeighborsView<ComputeSpace>& neighbors,
  //     const CumulativeProbabilityView<ComputeSpace>& 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 =
  //       mask_do_serch * static_cast<int>(neighbors.extent(1));
 
  //   std::size_t next =
  //       i_compartment * (1 - mask_do_serch) +
  //       mask_do_serch *
  //           neighbors(i_compartment, 0); // Default to the first neighbor
 
  //   for (int k_neighbor = 0; k_neighbor < max_neighbor - 1; ++k_neighbor)
  //   {
 
  //     // Get the cumulative probability range for the current neighbor
  //     const auto pi = cumulative_probability(i_compartment, k_neighbor);
  //     const auto pn = cumulative_probability(i_compartment, k_neighbor + 1);
 
  //     const auto cond =
  //         static_cast<int>(random_number <= pn && pi <= random_number);
 
  //     // Use of a Condition mask to avoid branching.
  //     next = cond * neighbors(i_compartment, k_neighbor + 1) + cond * next;
  //   }
 
  //   return next;
  // }
 
  struct TagMove
  {
  };
  struct TagLeave
  {
  };
 
  struct MoveFunctor
  {
    MoveFunctor() = default;
    MoveFunctor(MC::ParticlePositions p,
                MC::ParticleStatus _status,
                MC::DomainState<ComputeSpace> m,
                MC::pool_type _random_pool,
                MC::EventContainer _events,
                ProbeAutogeneratedBuffer _probes,
                MC::ParticleAges _ages,
                MC::ParticleSamples _random)
        : MoveFunctor(0,
                      std::move(p),
                      std::move(_status),
                      0,
                      std::move(m),
                      std::move(_random_pool),
                      std::move(_events),
                      std::move(_probes),
                      std::move(_ages),
                      std::move(_random),
                      false,
                      false) {};
 
    MoveFunctor(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,
                MC::ParticleSamples _random,
                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)),
          random(std::move(_random)), enable_move(b_move),
          enable_leave(b_leave) {};
 
    void
    update(const ComputeSpace& ex,
           double _d_t,
           std::size_t n_p,
           MC::DomainState<ComputeSpace> move_i,
           MC::ParticlePositions _positions,
           MC::ParticleStatus _status,
           MC::ParticleAges _ages,
           MC::ParticleSamples _random,
           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->random = std::move(_random);
 
      this->positions = std::move(_positions);
      this->status = std::move(_status);
      this->ages = std::move(_ages);
 
      Kokkos::fill_random(ex, random, random_pool, 0., 1.);
    }
 
    KOKKOS_INLINE_FUNCTION void
    operator()(
        TagMove /*tag*/,
        const Kokkos::TeamPolicy<ComputeSpace>::member_type& team_handle) const
    {
      GET_INDEX(n_particles);
      if (status(idx) != MC::Status::Idle) [[unlikely]]
      {
        return;
      }
 
      handle_move(idx);
    }
 
    KOKKOS_INLINE_FUNCTION void
    operator()(TagMove /*tag*/, const std::size_t& idx) const
    {
 
      if (status(idx) != MC::Status::Idle) [[unlikely]]
      {
        return;
      }
 
      handle_move(idx);
    }
 
    // KOKKOS_INLINE_FUNCTION void
    // operator()(TagLeave _tag,
    //            const Kokkos::TeamPolicy<ComputeSpace>::member_type& team,
    //            std::size_t& local_dead_count) const
    // {
    //   const auto work_stride = team.team_size() * team.league_size();
    //   const auto work_end = n_particles;
 
    //   std::size_t iwork =
    //       team.league_rank() * team.team_size() + team.team_rank();
 
    //   for (; iwork < work_end; iwork += work_stride)
    //   {
    //     ages(iwork, 0) += d_t;
    //     handle_exit(iwork, local_dead_count);
    //   }
    // }
    //
    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.liquid_volume, 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
    {
      // auto generator = random_pool.get_state();
      // const float rng1 = generator.frand(0., 1.);
      // const double rng2 = generator.drand(0., 1.);
      // random_pool.free_state(generator);
 
      const auto rng1 = static_cast<float>(random(idx, 0));
      const auto rng2 = static_cast<float>(random(idx, 1));
 
      const std::size_t i_current_compartment = positions(idx);
 
      KOKKOS_ASSERT(i_current_compartment < move.liquid_volume.extent(0));
 
      const bool mask_next = probability_leaving<void>(
          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));
 
      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
    {
      // const float random_number = random(idx, i);
 
      // SAMPLE_RANDOM_VARIABLES(random_pool,
      //                         const float random_number =
      //                             _generator_state_.frand(0., 1.);)
 
      // FIXME when move AND exit, take the same random number, is it really
      // important ?
      const auto random_number = static_cast<float>(random(idx, 0));
      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)
        {
          probes.set(ages(idx, 0));
        }
 
        // if (is_leaving && !probes.set(probe_value))
        // {
        //   Kokkos::printf("[Kernel]: PROBES OVERFLOW\r\n");
        // }
      }
      if constexpr (AutoGenerated::FlagCompileTime::enable_event_counter)
      {
        // if (is_leaving)
        // {
        //   events.incr<MC::EventType::Exit>();
        // }
        events.add<MC::EventType::Exit>(leave_mask);
      }
 
      dead_count += leave_mask;
      ages(idx, 0) = leave_mask * 0 + (1 - leave_mask) * ages(idx, 0);
      // status(idx) = is_leaving ? MC::Status::Exit : status(idx);
      status(idx) = static_cast<MC::Status>(
          static_cast<int>(MC::Status::Exit) * leave_mask
          + (1 - leave_mask) * static_cast<int>(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);
      }
    }
 
    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;
    // Kokkos::View<float**, Kokkos::LayoutLeft> random;
 
    MC::ParticleSamples random;
 
    bool enable_move{};
    bool enable_leave{};
  };
} // namespace Simulation::KernelInline
 
#endif