probe.hpp

#ifndef __CORE__PROBE_HPP__
#define __CORE__PROBE_HPP__
 
#include "Kokkos_Core_fwd.hpp"
#include <Kokkos_Core.hpp>
#include <biocma_cst_config.hpp>
#include <cstddef>
#include <cstdint>
#include <unordered_map>
// #include <type_traits>
 
// template <typename ViewType, typename ExecutionSpace>
// KOKKOS_INLINE_FUNCTION bool is_view_accessible(const ViewType& view)
// {
//   return Kokkos::SpaceAccessibility<typename ExecutionSpace::memory_space,
//                                     typename ViewType::memory_space>::accessible;
// }
 
namespace Simulation
{
 
  enum class ProbeType
  {
    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
 
  template <std::size_t buffer_size> class Probes
  {
  public:
    template <typename Space>
    using buffer_type = Kokkos::View<double[buffer_size], Kokkos::LayoutRight, Space>;
    void clear();
 
    [[nodiscard]] KOKKOS_INLINE_FUNCTION bool set(double val) const;
 
    [[nodiscard]] bool need_export() const noexcept;
 
    // [[nodiscard]] const double* raw_get() const;
 
    [[nodiscard]] std::span<const double> get() const;
 
    Probes();
 
  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>;
 
} // namespace Simulation
 
namespace Simulation
{
  template <std::size_t buffer_size> Probes<buffer_size>::Probes()
  {
    if (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();
  }
 
  template <std::size_t buffer_size>
  KOKKOS_INLINE_FUNCTION bool Probes<buffer_size>::set(double val) const
  {
    const auto i = (active)?Kokkos::atomic_fetch_inc(&internal_counter()):0;
 
    // Innaccessible from host space if use CUDA because of direct assignment
    if ((i < buffer_size) && active)
    {
      // this->buffer(i) = val;
      // if constexpr (is_view_accessible<decltype(buffer), Kokkos::DefaultExecutionSpace>())
      {
        this->buffer(i) = val;
      }
      // else
      // {
      //   // Kokkos::deep_copy(buffer(i), val);
      //   Kokkos::View<double, Kokkos::DefaultHostExecutionSpace> single_val("SingleValue");
      //   single_val() = val;
 
      //   Kokkos::deep_copy(buffer(i), single_val);
      // }
      return true;
    }
    return false;
  }
 
  template <std::size_t buffer_size> bool Probes<buffer_size>::need_export() const noexcept
  {
    return (internal_counter() >= buffer_size) && active;
  }
  template <std::size_t buffer_size> void Probes<buffer_size>::clear()
  {
    Kokkos::deep_copy(buffer, 0.);
    Kokkos::deep_copy(internal_counter, 0);
  }
 
  template <std::size_t buffer_size>
  [[nodiscard]] std::span<const double> Probes<buffer_size>::get() const
  {
    Kokkos::deep_copy(host_buffer, buffer);
    return {host_buffer.data(), std::min(buffer_size, internal_counter())};
  }
} // namespace Simulation
 
#endif