sched.h

#pragma once
 
#include "gaia/config/config.h"
 
#include <cstdint>
 
#include "gaia/cnt/darray.h"
#include "gaia/mem/smallblock_allocator.h"
#include "gaia/mt/jobcommon.h"
#include "gaia/mt/jobhandle.h"
#include "gaia/mt/threadpool.h"
 
namespace gaia {
  namespace ecs {
    enum class QueryExecType : uint32_t;
 
    struct SchedToken {
      uintptr_t value[2]{};
    };
 
    enum class SchedFlags : uint8_t {
      Default = 0,
      Background = 0x01
    };
 
    GAIA_NODISCARD inline bool sched_flags_has(SchedFlags flags, SchedFlags flag) {
      return ((uint8_t)flags & (uint8_t)flag) != 0U;
    }
 
    struct SchedTaskDesc {
      void* pCtx = nullptr;
      void (*invoke)(void* pCtx) = nullptr;
      QueryExecType execType{};
      SchedFlags flags = SchedFlags::Default;
    };
 
    struct SchedParDesc {
      void* pCtx = nullptr;
      void (*invoke)(void* pCtx, uint32_t idxStart, uint32_t idxEnd) = nullptr;
      uint32_t itemCount = 0;
      uint32_t groupSize = 0;
      QueryExecType execType{};
      SchedFlags flags = SchedFlags::Default;
    };
 
    struct Sched {
      void* pCtx = nullptr;
      SchedToken (*sched)(void* pCtx, const SchedTaskDesc* pDesc) = nullptr;
      SchedToken (*sched_par)(void* pCtx, const SchedParDesc* pDesc) = nullptr;
      SchedToken (*add)(void* pCtx, const SchedTaskDesc* pDesc) = nullptr;
      SchedToken (*add_par)(void* pCtx, const SchedParDesc* pDesc) = nullptr;
      void (*submit)(void* pCtx, SchedToken token) = nullptr;
      void (*dep)(void* pCtx, SchedToken tokenFirst, SchedToken tokenSecond) = nullptr;
      void (*wait)(void* pCtx, SchedToken token) = nullptr;
      void (*del)(void* pCtx, SchedToken token) = nullptr;
    };
 
    class SchedJob {
      Sched m_sched;
      SchedToken m_token{};
      void* m_pCleanupCtx = nullptr;
      void (*m_cleanup)(void* pCtx) = nullptr;
      bool m_valid = false;
      bool m_submitted = false;
      bool m_waited = false;
 
      void cleanup() {
        if (m_cleanup != nullptr) {
          m_cleanup(m_pCleanupCtx);
          m_cleanup = nullptr;
          m_pCleanupCtx = nullptr;
        }
      }
 
      GAIA_NODISCARD static bool same_sched(const Sched& a, const Sched& b) {
        return a.pCtx == b.pCtx && a.sched == b.sched && a.sched_par == b.sched_par && a.add == b.add &&
               a.add_par == b.add_par && a.submit == b.submit && a.dep == b.dep && a.wait == b.wait && a.del == b.del;
      }
 
    public:
      SchedJob() = default;
 
      SchedJob(Sched sched, SchedToken token, bool submitted, void* pCleanupCtx, void (*cleanup)(void* pCtx)):
          m_sched(sched), m_token(token), m_pCleanupCtx(pCleanupCtx), m_cleanup(cleanup), m_valid(true),
          m_submitted(submitted) {}
 
      ~SchedJob() {
        del();
      }
 
      SchedJob(const SchedJob&) = delete;
      SchedJob& operator=(const SchedJob&) = delete;
 
      SchedJob(SchedJob&& other) noexcept:
          m_sched(other.m_sched), m_token(other.m_token), m_pCleanupCtx(other.m_pCleanupCtx),
          m_cleanup(other.m_cleanup), m_valid(other.m_valid), m_submitted(other.m_submitted), m_waited(other.m_waited) {
        other.m_valid = false;
        other.m_cleanup = nullptr;
        other.m_pCleanupCtx = nullptr;
      }
 
      SchedJob& operator=(SchedJob&& other) noexcept {
        if (this == &other)
          return *this;
        del();
        m_sched = other.m_sched;
        m_token = other.m_token;
        m_pCleanupCtx = other.m_pCleanupCtx;
        m_cleanup = other.m_cleanup;
        m_valid = other.m_valid;
        m_submitted = other.m_submitted;
        m_waited = other.m_waited;
        other.m_valid = false;
        other.m_cleanup = nullptr;
        other.m_pCleanupCtx = nullptr;
        return *this;
      }
 
      GAIA_NODISCARD bool valid() const {
        return m_valid;
      }
 
      GAIA_NODISCARD SchedToken token() const {
        return m_token;
      }
 
      void submit();
      void dep(const SchedJob& jobFirst);
      void wait();
      void del();
    };
 
    namespace detail {
      inline mt::JobPriority exec_prio(QueryExecType execType) {
        // QueryExecType::ParallelEff is encoded as value 3. Keep the scheduler bridge independent
        // from the enum definition point so this header stays C-like and forward-declarable.
        return (uint32_t)execType == 3U ? mt::JobPriority::Low : mt::JobPriority::High;
      }
 
      inline bool sched_flags_background(SchedFlags flags) {
        return sched_flags_has(flags, SchedFlags::Background);
      }
 
      inline mt::JobCreationFlags job_creation_flags(SchedFlags flags) {
        uint8_t jobFlags = (uint8_t)mt::JobCreationFlags::ManualDelete;
        if (sched_flags_background(flags))
          jobFlags |= (uint8_t)mt::JobCreationFlags::Background;
        return (mt::JobCreationFlags)jobFlags;
      }
 
      enum class SchedTokenKind : uint32_t { None, Single, Parallel };
 
      struct SchedTokenDefData {
        cnt::darray<mt::JobHandle> handles;
        SchedTokenKind kind = SchedTokenKind::None;
        bool submitted = false;
 
        GAIA_USE_SMALLBLOCK(SchedTokenDefData)
      };
 
      inline SchedToken make_sched_token(SchedTokenDefData* pData) {
        SchedToken token{};
        token.value[0] = (uintptr_t)pData;
        return token;
      }
 
      inline SchedTokenDefData* sched_token_data(SchedToken token) {
        return reinterpret_cast<SchedTokenDefData*>(token.value[0]);
      }
 
      inline SchedToken add_one_def([[maybe_unused]] void* pCtx, const SchedTaskDesc* pDesc) {
        GAIA_ASSERT(pDesc != nullptr);
        GAIA_ASSERT(pDesc->invoke != nullptr);
        if (pDesc == nullptr || pDesc->invoke == nullptr)
          return {};
 
        auto* pData = new SchedTokenDefData();
        pData->kind = SchedTokenKind::Single;
        pData->handles.resize(1);
 
        mt::Job job;
        job.priority = exec_prio(pDesc->execType);
        job.flags = job_creation_flags(pDesc->flags);
        job.func = [pCtx = pDesc->pCtx, invoke = pDesc->invoke]() {
          invoke(pCtx);
        };
 
        pData->handles[0] = mt::ThreadPool::get().add(GAIA_MOV(job));
        return make_sched_token(pData);
      }
 
      inline SchedToken sched_one_def([[maybe_unused]] void* pCtx, const SchedTaskDesc* pDesc) {
        auto token = add_one_def(pCtx, pDesc);
        auto* pData = sched_token_data(token);
        if (pData != nullptr) {
          mt::ThreadPool::get().submit(pData->handles[0]);
          pData->submitted = true;
        }
        return token;
      }
 
      inline SchedToken add_par_def([[maybe_unused]] void* pCtx, const SchedParDesc* pDesc) {
        GAIA_ASSERT(pDesc != nullptr);
        GAIA_ASSERT(pDesc->invoke != nullptr);
        if (pDesc == nullptr || pDesc->invoke == nullptr || pDesc->itemCount == 0)
          return {};
 
        auto& tp = mt::ThreadPool::get();
        const auto prio = exec_prio(pDesc->execType);
        uint32_t groupSize = pDesc->groupSize;
        if (groupSize == 0) {
          const bool background = sched_flags_background(pDesc->flags);
          const auto workers =
              background ? core::get_max(1U, tp.background_workers()) : core::get_max(1U, tp.workers() + 1U);
          groupSize = (pDesc->itemCount + workers - 1) / workers;
          constexpr uint32_t maxUnitsOfWorkPerGroup = 8;
          groupSize = groupSize / maxUnitsOfWorkPerGroup;
          if (groupSize == 0)
            groupSize = 1;
        }
 
        const auto jobs = (pDesc->itemCount + groupSize - 1) / groupSize;
        auto* pData = new SchedTokenDefData();
        pData->kind = SchedTokenKind::Parallel;
        pData->handles.resize(jobs + 1);
 
        for (uint32_t jobIndex = 0; jobIndex < jobs; ++jobIndex) {
          const uint32_t idxStart = jobIndex * groupSize;
          const uint32_t idxEnd = core::get_min(idxStart + groupSize, pDesc->itemCount);
 
          mt::Job job;
          job.priority = prio;
          job.flags = job_creation_flags(pDesc->flags);
          job.func = [desc = *pDesc, idxStart, idxEnd]() {
            desc.invoke(desc.pCtx, idxStart, idxEnd);
          };
          pData->handles[jobIndex] = tp.add(GAIA_MOV(job));
        }
        {
          mt::Job syncJob;
          syncJob.priority = prio;
          syncJob.flags = job_creation_flags(pDesc->flags);
          pData->handles[jobs] = tp.add(GAIA_MOV(syncJob));
        }
        tp.dep(std::span(pData->handles.data(), jobs), pData->handles[jobs]);
        return make_sched_token(pData);
      }
 
      inline SchedToken sched_par_def([[maybe_unused]] void* pCtx, const SchedParDesc* pDesc) {
        auto token = add_par_def(pCtx, pDesc);
        auto* pData = sched_token_data(token);
        if (pData != nullptr) {
          mt::ThreadPool::get().submit(std::span(pData->handles.data(), pData->handles.size()));
          pData->submitted = true;
        }
        return token;
      }
 
      inline void sched_submit_def([[maybe_unused]] void* pCtx, SchedToken token) {
        auto* pData = sched_token_data(token);
        if (pData == nullptr || pData->submitted)
          return;
        auto& tp = mt::ThreadPool::get();
        tp.submit(std::span(pData->handles.data(), pData->handles.size()));
        pData->submitted = true;
      }
 
      inline void sched_dep_def([[maybe_unused]] void* pCtx, SchedToken tokenFirst, SchedToken tokenSecond) {
        auto* pFirst = sched_token_data(tokenFirst);
        auto* pSecond = sched_token_data(tokenSecond);
        if (pFirst == nullptr || pSecond == nullptr || pFirst->handles.empty() || pSecond->handles.empty())
          return;
        auto& tp = mt::ThreadPool::get();
        const auto firstDone = pFirst->handles.back();
        if (pSecond->kind == SchedTokenKind::Parallel && pSecond->handles.size() > 1) {
          const auto childCount = (uint32_t)pSecond->handles.size() - 1;
          for (uint32_t i = 0; i < childCount; ++i)
            tp.dep(firstDone, pSecond->handles[i]);
          return;
        }
        tp.dep(firstDone, pSecond->handles.back());
      }
 
      inline void sched_wait_def([[maybe_unused]] void* pCtx, SchedToken token) {
        auto* pData = sched_token_data(token);
        if (pData == nullptr || pData->handles.empty())
          return;
        mt::ThreadPool::get().wait(pData->handles.back());
      }
 
      inline void sched_del_def([[maybe_unused]] void* pCtx, [[maybe_unused]] SchedToken token) {
        auto* pData = sched_token_data(token);
        if (pData == nullptr)
          return;
        auto& tp = mt::ThreadPool::get();
        for (auto handle: pData->handles) {
          if (handle != mt::JobNull)
            tp.del(handle);
        }
        delete pData;
      }
    } // namespace detail
 
    GAIA_NODISCARD inline const Sched& sched_def() {
      static const Sched sched = [] {
        Sched b{};
        b.sched = &detail::sched_one_def;
        b.sched_par = &detail::sched_par_def;
        b.add = &detail::add_one_def;
        b.add_par = &detail::add_par_def;
        b.submit = &detail::sched_submit_def;
        b.dep = &detail::sched_dep_def;
        b.wait = &detail::sched_wait_def;
        b.del = &detail::sched_del_def;
        return b;
      }();
      return sched;
    }
 
    GAIA_NODISCARD inline const Sched& sched_resolve(const Sched& sched) {
      if (sched.sched == nullptr && sched.sched_par == nullptr && sched.add == nullptr && sched.add_par == nullptr &&
          sched.submit == nullptr && sched.dep == nullptr && sched.wait == nullptr && sched.del == nullptr)
        return sched_def();
      return sched;
    }
 
    GAIA_NODISCARD inline SchedToken sched_one(const Sched& sched, const SchedTaskDesc& desc) {
      const auto& resolved = sched_resolve(sched);
      if (resolved.sched != nullptr)
        return resolved.sched(resolved.pCtx, &desc);
      GAIA_ASSERT(resolved.add != nullptr);
      GAIA_ASSERT(resolved.submit != nullptr);
      const auto token = resolved.add != nullptr ? resolved.add(resolved.pCtx, &desc) : SchedToken{};
      if (resolved.submit != nullptr)
        resolved.submit(resolved.pCtx, token);
      return token;
    }
 
    GAIA_NODISCARD inline SchedJob
    sched_add(const Sched& sched, const SchedTaskDesc& desc, void* pCleanupCtx, void (*cleanup)(void* pCtx)) {
      const auto& resolved = sched_resolve(sched);
      if (resolved.add != nullptr) {
        if (resolved.submit == nullptr || resolved.wait == nullptr || resolved.del == nullptr) {
          GAIA_ASSERT(false);
          if (cleanup != nullptr)
            cleanup(pCleanupCtx);
          return {};
        }
        return SchedJob(resolved, resolved.add(resolved.pCtx, &desc), false, pCleanupCtx, cleanup);
      }
 
      if (resolved.sched == nullptr || resolved.wait == nullptr || resolved.del == nullptr) {
        GAIA_ASSERT(false);
        if (cleanup != nullptr)
          cleanup(pCleanupCtx);
        return {};
      }
      return SchedJob(resolved, resolved.sched(resolved.pCtx, &desc), true, pCleanupCtx, cleanup);
    }
 
    GAIA_NODISCARD inline SchedToken sched_par(const Sched& sched, const SchedParDesc& desc) {
      const auto& resolved = sched_resolve(sched);
      if (resolved.sched_par != nullptr)
        return resolved.sched_par(resolved.pCtx, &desc);
      GAIA_ASSERT(resolved.add_par != nullptr);
      GAIA_ASSERT(resolved.submit != nullptr);
      const auto token = resolved.add_par != nullptr ? resolved.add_par(resolved.pCtx, &desc) : SchedToken{};
      if (resolved.submit != nullptr)
        resolved.submit(resolved.pCtx, token);
      return token;
    }
 
    GAIA_NODISCARD inline SchedJob
    sched_add_par(const Sched& sched, const SchedParDesc& desc, void* pCleanupCtx, void (*cleanup)(void* pCtx)) {
      const auto& resolved = sched_resolve(sched);
      if (resolved.add_par != nullptr) {
        if (resolved.submit == nullptr || resolved.wait == nullptr || resolved.del == nullptr) {
          GAIA_ASSERT(false);
          if (cleanup != nullptr)
            cleanup(pCleanupCtx);
          return {};
        }
        return SchedJob(resolved, resolved.add_par(resolved.pCtx, &desc), false, pCleanupCtx, cleanup);
      }
 
      if (resolved.sched_par == nullptr || resolved.wait == nullptr || resolved.del == nullptr) {
        GAIA_ASSERT(false);
        if (cleanup != nullptr)
          cleanup(pCleanupCtx);
        return {};
      }
      return SchedJob(resolved, resolved.sched_par(resolved.pCtx, &desc), true, pCleanupCtx, cleanup);
    }
 
    inline void sched_submit(const Sched& sched, SchedToken token) {
      const auto& resolved = sched_resolve(sched);
      if (resolved.submit != nullptr)
        resolved.submit(resolved.pCtx, token);
    }
 
    inline void sched_dep(const Sched& sched, SchedToken tokenFirst, SchedToken tokenSecond) {
      const auto& resolved = sched_resolve(sched);
      if (resolved.dep != nullptr)
        resolved.dep(resolved.pCtx, tokenFirst, tokenSecond);
    }
 
    inline void sched_wait(const Sched& sched, SchedToken token) {
      const auto& resolved = sched_resolve(sched);
      if (resolved.wait != nullptr)
        resolved.wait(resolved.pCtx, token);
    }
 
    inline void sched_del(const Sched& sched, SchedToken token) {
      const auto& resolved = sched_resolve(sched);
      if (resolved.del != nullptr)
        resolved.del(resolved.pCtx, token);
    }
 
    inline void SchedJob::submit() {
      if (!m_valid || m_submitted)
        return;
      GAIA_ASSERT(m_sched.submit != nullptr);
      if (m_sched.submit == nullptr)
        return;
      sched_submit(m_sched, m_token);
      m_submitted = true;
    }
 
    inline void SchedJob::dep(const SchedJob& jobFirst) {
      if (!m_valid || !jobFirst.m_valid)
        return;
      GAIA_ASSERT(same_sched(m_sched, jobFirst.m_sched));
      GAIA_ASSERT(!m_submitted && !jobFirst.m_submitted);
      if (!same_sched(m_sched, jobFirst.m_sched) || m_submitted || jobFirst.m_submitted)
        return;
      sched_dep(m_sched, jobFirst.m_token, m_token);
    }
 
    inline void SchedJob::wait() {
      if (!m_valid || m_waited)
        return;
      GAIA_ASSERT(m_submitted);
      if (!m_submitted)
        return;
      sched_wait(m_sched, m_token);
      m_waited = true;
      cleanup();
    }
 
    inline void SchedJob::del() {
      if (!m_valid)
        return;
      if (m_submitted && !m_waited)
        wait();
      sched_del(m_sched, m_token);
      cleanup();
      m_valid = false;
      m_submitted = false;
      m_waited = false;
    }
  } // namespace ecs
} // namespace gaia