debug_cuda_scheduler.hpp

Go to the documentation of this file.

 
#pragma once
 
#include "cuda_task.hpp"
#include "debug_scheduler.hpp"
 
#include <uni20/backend/cuda/device.hpp>
#include <uni20/backend/cuda/runtime.hpp>
 
namespace uni20::async
{
 
class DebugCudaScheduler final : public DebugScheduler, public ICudaScheduler {
  public:
    explicit DebugCudaScheduler(DebugSchedulerOptions options = {})
        : DebugScheduler(options), default_device_(cuda::Device::get(0).ordinal()), device_count_(cuda::Device::count())
    {}
 
    explicit DebugCudaScheduler(cuda::Device default_device, DebugSchedulerOptions options = {})
        : DebugScheduler(options), default_device_(default_device.ordinal()), device_count_(cuda::Device::count())
    {}
 
    using DebugScheduler::schedule;
 
    void schedule(CudaTask&& task) override
    {
      TaskRegistry::record_task_scheduled(task.coroutine_handle());
      if (task.set_scheduler(this)) this->enqueue_task(std::move(task));
    }
 
    void schedule(CudaTask&& task, int device) override
    {
      auto const validated_device = cuda::Device::get(device);
      cuda_promise(task.handle()).bind_device(validated_device.ordinal());
      TaskRegistry::record_task_scheduled(task.coroutine_handle());
      if (task.set_scheduler(this)) this->enqueue_task(std::move(task));
    }
 
  private:
    bool accepts_route(TaskRoute route) const noexcept override
    {
      if (route.domain == TaskDomain::host) return !route.cuda_device;
      if (route.domain != TaskDomain::cuda) return false;
      int const device = route.cuda_device.value_or(default_device_);
      return device >= 0 && device < device_count_;
    }
 
    bool can_direct_transfer(TaskHandle from, TaskHandle to) const noexcept override
    {
      if (from.domain() != TaskDomain::cuda) return true;
      return this->activation_device(from) == this->activation_device(to);
    }
 
    void resume_task(BasicTask& task) override
    {
      if (task.handle().domain() != TaskDomain::cuda)
      {
        DebugScheduler::resume_task(task);
        return;
      }
 
      int const device = this->activation_device(task.handle());
      cuda::ScopedDevice device_scope(device);
#if UNI20_ASYNC_DEBUG
      DEBUG_CHECK_EQUAL(cuda::Device::current().ordinal(), device);
#endif
      DebugScheduler::resume_task(task);
    }
 
    [[nodiscard]] int activation_device(TaskHandle task) const noexcept
    {
      return cuda_promise(task).device().value_or(default_device_);
    }
 
    int default_device_;
    int device_count_;
};
 
} // namespace uni20::async