JobGraph.hpp

Go to the documentation of this file.

#pragma once
#include "Logging.hpp"
#include "ThreadPool.hpp"
#include <condition_variable>
#include <mutex>
namespace Foundation::Core
{
    struct JobHandle
    {
        static constexpr size_t kInvalid = ~static_cast<size_t>(0);
        size_t id{kInvalid};
        [[nodiscard]] bool Valid() const noexcept { return id != kInvalid; }
    };
 
    class JobGraph
    {
    public:
        enum class NodeKind
        {
            Worker, // single body, runs once on a pool worker (or main if no workers)
            ParallelFor, // body invoked per index, fanned out across the pool
            Main, // single body, runs on the Wait/Pump thread
            Barrier, // no body; completes when its producers complete
        };
 
    private:
        // Type-erased node body. Single-shot nodes use RunSingle; parallel-for nodes use RunIndex.
        struct IJobWork
        {
            virtual ~IJobWork() = default;
            virtual void RunSingle(size_t /*workerId*/) {}
            virtual void RunIndex(size_t /*index*/, size_t /*workerId*/) {}
        };
        template <typename Fn>
        struct SingleWork final : IJobWork
        {
            Fn fn;
            explicit SingleWork(Fn&& f) : fn(std::move(f)) {}
            void RunSingle(size_t workerId) override
            {
                if constexpr (std::is_invocable_v<Fn&, size_t>)
                    fn(workerId);
                else
                    fn();
            }
        };
        template <typename Fn>
        struct IndexWork final : IJobWork
        {
            Fn fn;
            explicit IndexWork(Fn&& f) : fn(std::move(f)) {}
            void RunIndex(size_t index, size_t workerId) override { ParallelForInvoke(fn, index, workerId); }
        };
 
        struct Node
        {
            NodeKind kind{NodeKind::Barrier};
            ExecutionPolicy policy{ExecutionPolicy::Par};
            UniquePtr<IJobWork> work{nullptr, StlDeleter<IJobWork>{nullptr}};
            size_t total{0}; // parallel-for element count
            StringView name;
            // Scheduling state guarded by JobGraph::mMutex.
            size_t pending{0}; // unmet producer dependencies
            bool finished{false};
            Vector<size_t> dependents;
            // Parallel-for hot state (touched lock-free by the fanned-out workers).
            Atomic<size_t> cursor{0};
            Atomic<size_t> remaining{0};
            explicit Node(Allocator* alloc) : dependents(alloc) {}
        };
 
        // Pool jobs that drive node bodies; defined in the .cpp (need the full JobGraph).
        struct WorkerJob;
        struct ParallelForJob;
 
        ThreadPool& mPool;
        Allocator* mAllocator;
        Vector<UniquePtr<Node>> mNodes;
        Mutex mMutex;
        CondVar mCond;
        Queue<size_t> mMainReady; // node ids ready to run on the Wait/Pump thread
        size_t mNodesRemaining{0};
        size_t mMainWorkerId{0};
        bool mSubmitted{false};
 
        size_t AddNode(NodeKind kind, StringView name, UniquePtr<IJobWork> work, size_t total,
                       ExecutionPolicy policy);
        void AddEdge(size_t producer, size_t consumer);
        void Schedule(size_t id);
        void DispatchWorker(size_t id);
        void DispatchParallelFor(size_t id);
        void EnqueueMain(size_t id);
        void RunMainNode(size_t id);
        void OnNodeFinished(size_t id);
 
    public:
        JobGraph(ThreadPool& pool, Allocator* allocator);
        JobGraph(JobGraph const&) = delete;
        JobGraph& operator=(JobGraph const&) = delete;
        ~JobGraph();
 
        template <typename Fn>
        JobHandle AddJob(StringView name, ExecutionPolicy policy, Fn&& fn)
        {
            auto work = ConstructUniqueBase<IJobWork, SingleWork<std::remove_reference_t<Fn>>>(
                mAllocator, std::forward<Fn>(fn));
            return {AddNode(NodeKind::Worker, name, std::move(work), 0, policy)};
        }
        template <typename Fn>
        JobHandle AddMain(StringView name, Fn&& fn)
        {
            auto work = ConstructUniqueBase<IJobWork, SingleWork<std::remove_reference_t<Fn>>>(
                mAllocator, std::forward<Fn>(fn));
            return {AddNode(NodeKind::Main, name, std::move(work), 0, ExecutionPolicy::Seq)};
        }
        template <typename Fn>
        JobHandle AddParallelFor(StringView name, ExecutionPolicy policy, size_t count, Fn&& fn)
        {
            auto work = ConstructUniqueBase<IJobWork, IndexWork<std::remove_reference_t<Fn>>>(
                mAllocator, std::forward<Fn>(fn));
            return {AddNode(NodeKind::ParallelFor, name, std::move(work), count, policy)};
        }
        template <typename It, typename Fn>
            requires std::random_access_iterator<It>
        JobHandle AddParallelFor(StringView name, ExecutionPolicy policy, It first, It last, Fn&& fn)
        {
            auto const count = last - first;
            size_t const total = count <= 0 ? 0 : static_cast<size_t>(count);
            return AddParallelFor(name, policy, total,
                [first, fn = std::forward<Fn>(fn)](size_t i, size_t worker) mutable
                { ParallelForInvoke(fn, first[static_cast<std::iter_difference_t<It>>(i)], worker); });
        }
        JobHandle AddBarrier(StringView name);
 
        template <typename... Producers>
        void DependsOn(JobHandle consumer, Producers... producers)
        {
            (AddEdge(producers.id, consumer.id), ...);
        }
 
        void Submit();
        void Wait(JobHandle target);
        void Wait() { Wait(JobHandle{}); }
        void PumpMainThread();
 
        [[nodiscard]] size_t MainWorkerId() const noexcept { return mMainWorkerId; }
    };
} // namespace Foundation::Core