map_completion.hh

Go to the documentation of this file.

#pragma once
 
// DiFfRG
#include <DiFfRG/common/kokkos.hh>
#include <DiFfRG/physics/integration/map_scheduler.hh>
 
// std
#include <cstring>
#include <exception>
#include <functional>
#include <type_traits>
#include <utility>
#include <vector>
 
namespace DiFfRG
{
  namespace internal
  {
    inline constexpr bool has_device_backend = !std::is_same_v<typename GPU_exec::memory_space, CPU_memory>;
  } // namespace internal
 
  class MapCompletion
  {
  public:
    struct PendingCopy {
      void *dst;
      const void *src;
      size_t bytes;
    };
 
    static void record(void *dst, const void *src, const size_t bytes)
    {
      pending().push_back(PendingCopy{dst, src, bytes});
    }
 
    static void record_work(std::function<void()> job) { deferred_work().push_back(std::move(job)); }
 
    static void flush()
    {
      // Host work first, and *before* the fence: the device is still chewing through whatever was
      // launched inside the scope, so this is the window in which host kernels are free. Drained
      // through a local so that a job which itself queues work cannot invalidate the iteration.
      auto &w = deferred_work();
      while (!w.empty()) {
        std::vector<std::function<void()>> batch;
        batch.swap(w);
        for (auto &job : batch)
          job();
      }
 
      Kokkos::fence("DiFfRG::MapCompletion::flush");
      auto &p = pending();
      for (const auto &c : p)
        std::memcpy(c.dst, c.src, c.bytes);
      p.clear();
      MapScheduler::instance().complete();
    }
 
    static void discard(const char *reason)
    {
      // Dropped without running: the destinations may already be gone, and an unrun job is exactly
      // as abandoned as an unlanded copy.
      deferred_work().clear();
      pending().clear();
      MapScheduler::instance().poison(reason);
    }
 
    static bool has_pending(const void *src)
    {
      for (const auto &c : pending())
        if (c.src == src) return true;
      return false;
    }
 
    static bool deferral_enabled() { return deferral(); }
    static void set_deferral(const bool enabled)
    {
      deferral() = enabled;
      // The scheduler splits a map that is flushed on return more aggressively than one inside a
      // batch, because the ranks it leaves out have nothing else to pick up. See
      // MapScheduler::set_batched().
      MapScheduler::instance().set_batched(enabled);
    }
 
  private:
    static std::vector<PendingCopy> &pending()
    {
      static std::vector<PendingCopy> p;
      return p;
    }
    static std::vector<std::function<void()>> &deferred_work()
    {
      static std::vector<std::function<void()>> w;
      return w;
    }
    static bool &deferral()
    {
      static bool enabled = false;
      return enabled;
    }
  };
 
  inline void flush_maps() { MapCompletion::flush(); }
 
  class DeferredMaps
  {
  public:
    DeferredMaps() { MapCompletion::set_deferral(true); }
    ~DeferredMaps()
    {
      MapCompletion::set_deferral(false);
      // Leaving via an exception means this rank abandons the batch while the others are still
      // filling theirs. Running the collective here would deadlock; poisoning the scheduler turns
      // the next flush into a diagnosable abort instead. Under MPI an exception thrown out of a
      // flow block is not recoverable anyway -- the ranks have already diverged.
      if (std::uncaught_exceptions() > 0) {
        MapCompletion::discard("an exception was thrown inside a DeferredMaps scope");
        return;
      }
      MapCompletion::flush();
    }
    DeferredMaps(const DeferredMaps &) = delete;
    DeferredMaps &operator=(const DeferredMaps &) = delete;
  };
} // namespace DiFfRG