tbb_numa_scheduler.hpp

Go to the documentation of this file.

#pragma once
 
#include "scheduler.hpp"
#include "tbb_scheduler.hpp"
 
#include <uni20/common/trace.hpp>
 
#include <atomic>
#include <fmt/core.h>
#include <memory>
#include <mutex>
#include <optional>
#include <unordered_map>
#include <vector>
 
#include <oneapi/tbb/info.h>
 
namespace uni20::async
{
 
class TbbNumaScheduler final : public IAsyncScheduler {
  public:
    TbbNumaScheduler() : TbbNumaScheduler(oneapi::tbb::info::numa_nodes()) {}
 
    explicit TbbNumaScheduler(std::vector<int> nodes) : numa_nodes_(std::move(nodes))
    {
      if (numa_nodes_.empty())
      {
        numa_nodes_.push_back(-1);
      }
 
      arenas_.reserve(numa_nodes_.size());
      scheduled_counts_.assign(numa_nodes_.size(), 0);
 
      for (std::size_t i = 0; i < numa_nodes_.size(); ++i)
      {
        auto node = numa_nodes_[i];
        oneapi::tbb::task_arena::constraints constraints;
        constraints.set_numa_id(node);
        auto scheduler = std::make_unique<TbbScheduler>(constraints);
        node_to_index_.emplace(node, arenas_.size());
        arenas_.push_back(Arena{node, std::move(scheduler)});
      }
 
      fmt::print("[uni20] TbbNumaScheduler: initialized {} NUMA nodes\n", arenas_.size());
    }
 
    void schedule(AsyncTask&& task) override
    {
      TaskRegistry::record_task_scheduled(task.coroutine_handle());
      int target = task.preferred_numa_node().value_or(this->select_next_numa_node());
      this->schedule_on_node(std::move(task), target);
    }
 
    void schedule(AsyncTask&& task, int numa_node)
    {
      TaskRegistry::record_task_scheduled(task.coroutine_handle());
      this->schedule_on_node(std::move(task), numa_node);
    }
 
    void pause() override
    {
      for (auto& arena : arenas_)
      {
        arena.scheduler->pause();
      }
    }
 
    void resume() override
    {
      for (auto& arena : arenas_)
      {
        arena.scheduler->resume();
      }
    }
 
    void run_all()
    {
      for (auto& arena : arenas_)
      {
        arena.scheduler->run_all();
      }
    }
 
    [[nodiscard]] const std::vector<int>& numa_nodes() const noexcept { return numa_nodes_; }
 
    [[nodiscard]] std::size_t scheduled_count_for(int numa_node) const noexcept
    {
      CHECK(!arenas_.empty(), "TbbNumaScheduler must retain at least one arena");
      auto idx = this->index_for_node(numa_node);
      std::lock_guard<std::mutex> lock(counts_mutex_);
      return scheduled_counts_[idx];
    }
 
  private:
    void execute_batch_impl(LightweightTaskBatch const& batch) override
    {
      CHECK(!arenas_.empty(), "TbbNumaScheduler must retain at least one arena");
      arenas_[this->index_for_node(this->select_next_numa_node())].scheduler->execute_batch(
          batch.size(), [&batch](std::size_t index) { batch(index); });
    }
 
    bool accepts_route(TaskRoute route) const noexcept override
    {
      return route.domain == TaskDomain::host && !route.cuda_device;
    }
 
    // TbbNumaScheduler itself is a router; direct execution belongs to one of
    // its child arena schedulers.
    bool can_direct_transfer(TaskHandle, TaskHandle) const noexcept override { return false; }
 
    void reschedule(BasicTask&& task) override
    {
      if (auto preferred = task.preferred_numa_node())
      {
        this->reschedule_on_node(std::move(task), *preferred);
      }
      else
      {
        this->reschedule_on_node(std::move(task), this->select_next_numa_node());
      }
    }
 
    struct Arena
    {
        int numa_node;
        std::unique_ptr<TbbScheduler> scheduler;
    };
 
    [[nodiscard]] std::size_t index_for_node(int numa_node) const noexcept
    {
      if (auto it = node_to_index_.find(numa_node); it != node_to_index_.end())
      {
        return it->second;
      }
      return 0;
    }
 
    [[nodiscard]] int select_next_numa_node() noexcept
    {
      CHECK(!arenas_.empty(), "TbbNumaScheduler must retain at least one arena");
      auto index = next_index_.fetch_add(1, std::memory_order_relaxed) % arenas_.size();
      return arenas_[index].numa_node;
    }
 
    void schedule_on_node(AsyncTask&& task, int numa_node)
    {
      CHECK(!arenas_.empty(), "TbbNumaScheduler must retain at least one arena");
      auto index = this->index_for_node(numa_node);
      auto& arena = arenas_[index];
      auto actual_node = arena.numa_node;
      task.set_preferred_numa_node(actual_node);
      {
        std::lock_guard<std::mutex> lock(counts_mutex_);
        ++scheduled_counts_[index];
      }
      arena.scheduler->schedule(std::move(task));
    }
 
    void reschedule_on_node(BasicTask&& task, int numa_node)
    {
      CHECK(!arenas_.empty(), "TbbNumaScheduler must retain at least one arena");
      auto index = this->index_for_node(numa_node);
      auto& arena = arenas_[index];
      auto actual_node = arena.numa_node;
      task.set_preferred_numa_node(actual_node);
      {
        std::lock_guard<std::mutex> lock(counts_mutex_);
        ++scheduled_counts_[index];
      }
      arena.scheduler->reschedule(std::move(task));
    }
 
    std::vector<int> numa_nodes_;
    std::vector<Arena> arenas_;
    std::unordered_map<int, std::size_t> node_to_index_;
    std::vector<std::size_t> scheduled_counts_;
    mutable std::mutex counts_mutex_;
    std::atomic<std::size_t> next_index_{0};
};
 
} // namespace uni20::async