debug_scheduler.hpp

Go to the documentation of this file.

 
 
#pragma once
 
#include "async.hpp"
#include "cuda_task.hpp"
#include "scheduler.hpp"
#include "task_registry.hpp"
#include <algorithm>
#include <cstdint>
#include <numeric>
#include <random>
#include <uni20/common/display.hpp>
#include <utility>
#include <vector>
 
namespace uni20::async
{
 
namespace detail
{
inline void service_task_registry_debug_requests()
{
#if UNI20_DEBUG_ASYNC_TASKS
  TaskRegistry::service_debug_requests();
#endif
}
 
inline void dump_deadlock_graphviz_snapshot()
{
#if UNI20_DEBUG_ASYNC_TASKS
  auto const path = TaskRegistry::default_graphviz_dump_path();
  if (TaskRegistry::dump_graphviz_file_best_effort(path))
  {
    display::info("Async Graphviz DAG snapshot written to {}", path);
  }
#endif
}
} // namespace detail
 
enum class DebugSchedulerOrder
{
  fifo,
  reverse,
  random,
};
 
struct DebugSchedulerOptions
{
    DebugSchedulerOrder order = DebugSchedulerOrder::reverse;
    std::uint64_t random_seed = 0;
};
 
class DebugScheduler : public IAsyncScheduler {
  public:
    explicit DebugScheduler(DebugSchedulerOptions options = {}) : options_(options), random_engine_(options.random_seed)
    {}
 
    void schedule(AsyncTask&& task) override
    {
      TRACE_MODULE(ASYNC, "Scheduling a task", &task, task.coroutine_handle());
      TaskRegistry::record_task_scheduled(task.coroutine_handle());
      if (task.set_scheduler(this))
      {
        this->enqueue_task(std::move(task));
      }
    }
 
    ~DebugScheduler()
    {
      TRACE_MODULE(ASYNC, "~DebugScheduler", Handles_.size());
 
      if (std::uncaught_exceptions() > 0)
      {
        while (!Handles_.empty())
        {
          Handles_.back().abandon_leak();
          Handles_.pop_back();
        }
      }
    }
 
    [[nodiscard]] bool can_run() const noexcept { return !Blocked_ && !Handles_.empty(); }
 
    void run();
 
    void run_all();
 
    void pause() override { Blocked_ = true; }
 
    void resume() override { Blocked_ = false; }
 
    void help_while_waiting(const WaitPredicate& is_ready) override
    {
      if (is_ready())
      {
        return;
      }
 
      if (Blocked_ || Handles_.empty())
      {
        detail::dump_deadlock_graphviz_snapshot();
        TaskRegistry::dump();
        PANIC("**DEADLOCK** get_wait object is not available but there are no runnable tasks!");
      }
      this->run();
    }
 
    [[nodiscard]] bool done() const noexcept { return Handles_.empty(); }
 
    [[nodiscard]] DebugSchedulerOptions const& options() const noexcept { return options_; }
 
  private:
    void execute_batch_impl(LightweightTaskBatch const& batch) override
    {
      std::vector<std::size_t> indices(batch.size());
      std::iota(indices.begin(), indices.end(), std::size_t{0});
      this->order_batch(indices);
      for (std::size_t const index : indices)
        batch(index);
    }
 
    bool accepts_route(TaskRoute route) const noexcept override
    {
      return route.domain == TaskDomain::host && !route.cuda_device;
    }
 
    bool can_direct_transfer(TaskHandle, TaskHandle) const noexcept override { return true; }
 
    // Internal resubmission
    void reschedule(BasicTask&& task) override
    {
      TRACE_MODULE(ASYNC, "Rescheduling a task", &task, task.coroutine_handle());
      // Assume sched_ is already set
      this->enqueue_task(std::move(task));
    }
 
    bool Blocked_ = false;
 
    std::vector<BasicTask> Handles_;
    DebugSchedulerOptions options_;
    std::mt19937_64 random_engine_;
 
    template <class T> void order_batch(std::vector<T>& batch)
    {
      switch (options_.order)
      {
        case DebugSchedulerOrder::fifo:
          return;
        case DebugSchedulerOrder::reverse:
          std::reverse(batch.begin(), batch.end());
          return;
        case DebugSchedulerOrder::random:
          // Keep the permutation stable across standard-library implementations.
          for (std::size_t remaining = batch.size(); remaining > 1; --remaining)
          {
            auto const selected = static_cast<std::size_t>(random_engine_() % remaining);
            std::swap(batch[remaining - 1], batch[selected]);
          }
          return;
      }
      PANIC("invalid DebugScheduler ordering policy", static_cast<int>(options_.order));
    }
 
  protected:
    void enqueue_task(BasicTask&& task) { Handles_.push_back(std::move(task)); }
 
    virtual void resume_task(BasicTask& task) { task.resume(); }
};
 
namespace detail
{
inline DebugScheduler DefaultScheduler;
inline IAsyncScheduler* global_scheduler = &DefaultScheduler;
} // namespace detail
 
inline void set_global_scheduler(IAsyncScheduler* sched) { detail::global_scheduler = sched; }
 
inline IAsyncScheduler* get_global_scheduler() { return detail::global_scheduler; }
 
inline void reset_global_scheduler() { detail::global_scheduler = &detail::DefaultScheduler; }
 
// ScopedScheduler is useful for testing; set the scheduler for the lifetime of a block
class ScopedScheduler {
  public:
    explicit ScopedScheduler(IAsyncScheduler* sched)
    {
      old_ = get_global_scheduler();
      set_global_scheduler(sched);
    }
    ~ScopedScheduler() { set_global_scheduler(old_); }
 
  private:
    IAsyncScheduler* old_;
};
 
inline void schedule(AsyncTask&& task) { get_global_scheduler()->schedule(std::move(task)); }
 
template <class Function>
  requires std::invocable<Function&, std::size_t>
void execute_batch(std::size_t size, Function&& function)
{
  get_global_scheduler()->execute_batch(size, std::forward<Function>(function));
}
 
inline void schedule(CudaTask&& task)
{
  auto* scheduler = dynamic_cast<ICudaScheduler*>(get_global_scheduler());
  CHECK(scheduler != nullptr, "the global scheduler does not accept CUDA tasks");
  scheduler->schedule(std::move(task));
}
 
inline void schedule(CudaTask&& task, int device)
{
  auto* scheduler = dynamic_cast<ICudaScheduler*>(get_global_scheduler());
  CHECK(scheduler != nullptr, "the global scheduler does not accept CUDA tasks");
  scheduler->schedule(std::move(task), device);
}
 
template <typename T> void EpochContextReader<T>::wait() const
{
  auto* sched = get_global_scheduler();
  if (!this->ready())
  {
    CHECK(sched);
    auto epoch = epoch_;
    sched->wait_for(IAsyncScheduler::WaitRequest{
        .is_ready = [epoch] { return epoch->reader_ready(); },
        .notify_when_ready =
            [epoch](IAsyncScheduler::WaitWakeup notify) { epoch->reader_notify_when_ready(std::move(notify)); },
    });
  }
}
 
template <typename T> void EpochContextReader<T>::wait(IAsyncScheduler& sched) const
{
  if (!this->ready())
  {
    auto epoch = epoch_;
    sched.wait_for(IAsyncScheduler::WaitRequest{
        .is_ready = [epoch] { return epoch->reader_ready(); },
        .notify_when_ready =
            [epoch](IAsyncScheduler::WaitWakeup notify) { epoch->reader_notify_when_ready(std::move(notify)); },
    });
  }
}
 
template <typename T> T const& EpochContextReader<T>::get_wait() const
{
  this->wait();
  return this->data();
}
 
template <typename T> T const& EpochContextReader<T>::get_wait(IAsyncScheduler& sched) const
{
  this->wait(sched);
  return this->data();
}
 
template <typename T> T&& EpochContextWriter<T>::move_from_wait()
{
  auto* sched = get_global_scheduler();
  if (!this->ready())
  {
    CHECK(sched);
    auto epoch = this->epoch_context_shared();
    sched->wait_for(IAsyncScheduler::WaitRequest{
        .is_ready = [this] { return this->ready(); },
        .notify_when_ready =
            [epoch](IAsyncScheduler::WaitWakeup notify) { epoch->writer_notify_when_ready(std::move(notify)); },
    });
  }
  return std::move(this->data());
}
 
//-----------------------------------------------------------------------------
// Inline definitions
//-----------------------------------------------------------------------------
 
inline void DebugScheduler::run()
{
  DEBUG_TRACE_MODULE(ASYNC, "DebugScheduler::run");
  detail::service_task_registry_debug_requests();
  if (Blocked_)
  {
    DEBUG_TRACE_MODULE(ASYNC, "run() on a blocked DebugQueue: doing nothing");
    return;
  }
  TRACE_MODULE(ASYNC, "Got some coroutines to resume", Handles_.size());
 
  std::vector<BasicTask> H;
  std::swap(H, Handles_);
  this->order_batch(H);
  for (auto&& h : H)
  {
    TRACE_MODULE(ASYNC, "resuming coroutine...", &h, h.coroutine_handle());
    this->resume_task(h);
    CHECK(!h);
    TRACE_MODULE(ASYNC, "here", &h, Handles_.size());
  }
  detail::service_task_registry_debug_requests();
}
 
inline void DebugScheduler::run_all()
{
  DEBUG_TRACE_MODULE(ASYNC, "DebugScheduler::run_all");
  detail::service_task_registry_debug_requests();
  if (Blocked_)
  {
    DEBUG_TRACE_MODULE(ASYNC, "run() on a blocked DebugQueue: doing nothing");
    return;
  }
  while (!done())
  {
    run();
  }
  detail::service_task_registry_debug_requests();
}
 
} // namespace uni20::async