work_stealing_executor.hpp

Go to the documentation of this file.

#pragma once

 
#include <cstdint>
#include <mutex>
#include <string_view>
#include <vector>
 
#include "orrery/backend/cpu_topology.hpp"
#include "orrery/backend/executor.hpp"
#include "orrery/backend/partition.hpp"
#include "orrery/backend/thread_pool.hpp"
#include "orrery/backend/worker_statistics.hpp"
#include "orrery/core/aligned_allocator.hpp"
#include "orrery/core/types.hpp"
 
namespace orrery::backend {

class WorkStealingExecutor final : public Executor {
public:
    static constexpr core::Index kChunksPerWorker = 16;
 
    explicit WorkStealingExecutor(unsigned worker_count = 0,
                                  ThreadPool::Affinity affinity = ThreadPool::Affinity::kUnpinned);
 
    void run(core::Index count, RangeTask task) override;
 
    [[nodiscard]] std::string_view name() const noexcept override { return "work-stealing"; }
 
    [[nodiscard]] unsigned worker_count() const noexcept override { return pool_.worker_count(); }
 
    [[nodiscard]] CoreClass worker_core_class(unsigned worker) const noexcept override {
        return pool_.core_class(worker);
    }
 
    [[nodiscard]] ExecutorStatistics statistics() const override;
 
    void reset_statistics() noexcept override;
 
private:
    struct alignas(core::kCacheLineBytes) WorkRange {
        std::mutex mutex;

        core::Index next{};

        core::Index end{};
    };

    void drain(unsigned worker) noexcept;

    [[nodiscard]] bool claim_front(unsigned worker, IndexRange& chunk) noexcept;

    [[nodiscard]] bool claim_back(unsigned victim, IndexRange& chunk) noexcept;

    [[nodiscard]] bool steal(unsigned thief, IndexRange& chunk) noexcept;
 
    ThreadPool pool_;
 
    std::vector<PaddedWorkerStatistics> workers_;
    std::vector<Duration> busy_at_region_start_;

    std::vector<WorkRange> ranges_;

    const RangeTask* task_{nullptr};

    core::Index grain_{kPartitionGrain};
 
    Duration elapsed_{};
    std::uint64_t regions_{};
};
 
} // namespace orrery::backend