probe.hpp

#ifndef __CORE__PROBE_HPP__
#define __CORE__PROBE_HPP__
 
#include <Kokkos_Core.hpp>
#include <biocma_cst_config.hpp>
#include <cstddef>
#include <cstdint>
#include <span>
#include <unordered_map>
 
namespace Simulation
{
 
  enum class ProbeType : unsigned
  {
    LeavingTime = 0,
    DivisionTime
  };
 
  // The name are namely used during export
  static const std::array<std::string, 2> map_probe_name
      = { "LeavingTime", "DivisionTime" };
 
  // TODO: keep thi structure as a bulkstorage for fast compute but consider
  // sorting into bins when exporting it to avoid unnecessary amount of data
  //  Unordered map could be adapted
 
  /*
  TODO, it seems that after some calculation (example cstr 0d) size(division
  probes) = number of tally+1 where the last probe is 0 where size(leaving time)
  == tally Is it a bug ?
  */

  template <std::size_t buffer_size> class Probes
  {
  public:
    template <typename Space>
    using buffer_type
        = Kokkos::View<double[buffer_size], Kokkos::LayoutRight, Space>;
 
    // WARNING : marked  const but mutates state
    void
    clear() const
    {
      if constexpr (AutoGenerated::FlagCompileTime::use_probe)
      {
        if (active)
        {
          Kokkos::deep_copy(buffer, 0.);
          Kokkos::deep_copy(internal_counter, 0);
        }
      }
    }
 
    template <typename Space>
    [[nodiscard]] KOKKOS_INLINE_FUNCTION bool
    set(double val) const
    {
      static_assert(Kokkos::SpaceAccessibility<
                    Space,
                    Kokkos::DefaultExecutionSpace::memory_space>::accessible);
      //Strategy here: for the moment kernel implementation do not allow to have use_probe==true and active ==false.
      // Then check active for safetiness but useless
      if constexpr (AutoGenerated::FlagCompileTime::use_probe)
      {
        if (const auto i = Kokkos::atomic_fetch_inc(&internal_counter());
            (i < buffer_size) && active)
        {
          this->buffer(i) = val;
          return true;
        }
        return false;
      }
      //Return true if probe is not activate indicates no error to caller
      return true;
    }
 
    [[nodiscard]] bool
    need_export() const noexcept
    {
      if constexpr (AutoGenerated::FlagCompileTime::use_probe)
      {
        return active && (internal_counter() >= buffer_size);
      }
      return false; // As the probe never need to be exported, there is no error
    }
 
    [[nodiscard]] std::span<const double>
    get() const
    {
      if constexpr (AutoGenerated::FlagCompileTime::use_probe)
      {
        if (!active)
        {
          // return empty span is a valid return ,
          //  use {} constructor to comply clang and gcc
          return {};
        }
        // TODO this should be alsway internal_counter(), but in case
        // internal counter > buffersize (error in race condition) size
        // never overrun buffer_size
        const auto return_size = std::min(buffer_size, internal_counter());
        assert(return_size <= buffer_size);
        /*
        // Safe: returning a span of size return_size is ok; the span may be
        smaller than the full host_buffer
        */
        Kokkos::deep_copy(host_buffer, buffer);
        return { host_buffer.data(), return_size };
      }
 
      return {}; // Also valid, probe not activated returns empty span
    }
 
    Probes()
    {
      if constexpr (AutoGenerated::FlagCompileTime::use_probe)
      {
        active = true;
        buffer = buffer_type<Kokkos::DefaultExecutionSpace>("probe_buffer");
        internal_counter = Kokkos::View<uint64_t, Kokkos::SharedSpace>("i_c");
        host_buffer
            = buffer_type<Kokkos::DefaultHostExecutionSpace>("host_buffer");
 
        clear();
      }
      active = false;
    }
 
  private:
    bool active{};
    // NOLINTBEGIN
    buffer_type<Kokkos::DefaultExecutionSpace> buffer;
    Kokkos::View<uint64_t, Kokkos::SharedSpace> internal_counter;
    buffer_type<Kokkos::DefaultHostExecutionSpace> host_buffer;
    // NOLINTEND
  };
 
  using ProbeAutogeneratedBuffer = Probes<AutoGenerated::probe_buffer_size>;
 
  using MapProbes = std::unordered_map<ProbeType, ProbeAutogeneratedBuffer>;
 
  using ProbeIterator = MapProbes::const_iterator;
 
} // namespace Simulation
#endif