map_scheduler.hh

Go to the documentation of this file.

#pragma once
 
// DiFfRG
#include <DiFfRG/common/kokkos.hh>
#include <DiFfRG/common/mpi.hh>
 
// std
#include <array>
#include <cstddef>
#include <cstdint>
#include <type_traits>
#include <vector>
 
namespace DiFfRG
{
  namespace internal
  {
    double device_fill_threshold();
    double host_fill_threshold();
 
    inline constexpr double host_grain = 1024.;
 
    inline constexpr double launch_threshold = 1024.;
  } // namespace internal
 
  enum class MapResource : int { device = 0, host = 1 };
  inline constexpr int n_map_resources = 2;
 
  inline const char *to_string(const MapResource r) { return r == MapResource::host ? "host" : "device"; }
 
  struct MapTarget {
    MapResource resource;
    double fill_threshold;
  };
 
  template <typename ExecutionSpace> inline MapTarget map_target()
  {
    if constexpr (std::is_same_v<typename ExecutionSpace::memory_space, CPU_memory>)
      return MapTarget{MapResource::host, internal::host_fill_threshold()};
    else
      return MapTarget{MapResource::device, internal::device_fill_threshold()};
  }
 
  template <typename ExecutionSpace> inline double map_fill_threshold()
  {
    return map_target<ExecutionSpace>().fill_threshold;
  }
 
  struct MapSlice {
    size_t offset = 0;
    size_t count = 0;
 
    bool owns_all(const size_t grid_size) const { return offset == 0 && count == grid_size; }
  };
 
  class NoMapsHere
  {
  public:
    NoMapsHere();
    ~NoMapsHere();
    NoMapsHere(const NoMapsHere &) = delete;
    NoMapsHere &operator=(const NoMapsHere &) = delete;
 
    static bool active();
  };
 
  class MapScheduler
  {
  public:
    static MapScheduler &instance();
 
    bool active() const { return m_n_ranks > 1; }
 
    uint rank() const { return m_rank; }
    uint n_ranks() const { return m_n_ranks; }
    MPI_Comm comm() const { return m_comm; }
 
    MapSlice schedule(size_t integrator_id, void *dest, size_t elem_size, size_t grid_size, size_t quadrature_volume,
                      bool splittable, MapTarget target);
 
    bool plan_contains(size_t integrator_id) const;
 
    bool has_open_plan() const { return !m_plan.empty(); }
 
    void complete();
 
    void poison(const char *reason);
 
    void set_batched(const bool batched) { m_batched = batched; }
 
    void set_quantum(double quantum);
    double quantum() const { return m_quantum_override; }
    void set_verbose(bool verbose) { m_verbose = verbose; }
 
  protected:
    MapScheduler();
 
    MapScheduler(uint rank, uint n_ranks);
 
    bool prepare_batch();
    void finish_batch();
 
    // Grow-only scratch, so a flush does not allocate. Protected rather than private so the
    // multi-rank simulation in the tests can play the part of the network.
    std::vector<char> m_send;
    std::vector<char> m_recv;
    std::vector<int> m_counts;
    std::vector<int> m_displs;
 
  private:
    struct Entry {
      size_t integrator_id;
      void *dest;
      size_t elem_size;
      size_t grid_size;
      MapResource resource;
      std::vector<uint> owners;
    };
 
    static size_t part(size_t grid_size, size_t r, size_t j) { return (j * grid_size) / r; }
 
    std::vector<uint> least_loaded(size_t r, MapResource resource) const;
    void reset_load();
    uint64_t plan_hash() const;
    void log_plan() const;
 
    MPI_Comm m_comm;
    uint m_rank = 0;
    uint m_n_ranks = 1;
 
    double m_quantum_override = 0.;
    bool m_quantum_pinned = false;
    bool m_batched = false;
    size_t m_batch_index = 0;
    size_t m_verify_every = 64;
    static constexpr size_t verify_head = 8;
    bool m_verbose = false;
    bool m_logged = false;
    const char *m_poisoned = nullptr;
 
    std::vector<Entry> m_plan;
    std::array<std::vector<double>, n_map_resources> m_load;
    bool m_simulated = false;
  };
} // namespace DiFfRG