jobmanager.h

#pragma once
 
#include "gaia/config/config.h"
#include "gaia/config/profiler.h"
 
#include <atomic>
#include <cinttypes>
// TODO: Currently necessary due to std::function. Replace them!
#include <functional>
 
#include "gaia/cnt/ilist.h"
#include "gaia/core/span.h"
#include "gaia/core/utility.h"
#include "gaia/mem/mem_alloc.h"
#include "gaia/mt/jobcommon.h"
#include "gaia/mt/jobhandle.h"
 
#define GAIA_LOG_JOB_STATES 0
 
namespace gaia {
  namespace mt {
    enum JobState : uint32_t {
      DEP_BITS_START = 0,
      DEP_BITS = 27,
      DEP_BITS_MASK = (uint32_t)((1u << DEP_BITS) - 1),
 
      STATE_BITS_START = DEP_BITS_START + DEP_BITS,
      STATE_BITS = 4,
      STATE_BITS_MASK = (uint32_t)(((1u << STATE_BITS) - 1) << STATE_BITS_START),
 
      // STATE
 
      Submitted = 0x01 << STATE_BITS_START,
      Processing = 0x02 << STATE_BITS_START,
      Executing = 0x03 << STATE_BITS_START,
      Done = 0x04 << STATE_BITS_START,
      Released = 0x05 << STATE_BITS_START
    };
 
    struct JobContainer;
 
    namespace detail {
      inline void signal_edge(JobContainer& jobData);
    }
 
    struct JobEdges {
      union {
        JobHandle dep;
        JobHandle* pDeps;
      };
      uint32_t depCnt;
    };
 
    struct JobContainer: cnt::ilist_item {
      std::atomic_uint32_t state;
      JobPriority prio;
      JobCreationFlags flags;
      JobEdges edges;
      std::function<void()> func;
 
      JobContainer() = default;
      ~JobContainer() = default;
 
      JobContainer(const JobContainer& other): cnt::ilist_item(other) {
        state = other.state.load();
        prio = other.prio;
        flags = other.flags;
        edges = other.edges;
        func = other.func;
      }
      JobContainer& operator=(const JobContainer& other) {
        GAIA_ASSERT(core::addressof(other) != this);
        cnt::ilist_item::operator=(other);
        state = other.state.load();
        prio = other.prio;
        flags = other.flags;
        edges = other.edges;
        func = other.func;
        return *this;
      }
 
      JobContainer(JobContainer&& other): cnt::ilist_item(GAIA_MOV(other)) {
        state = other.state.load();
        prio = other.prio;
        flags = other.flags;
        func = GAIA_MOV(other.func);
 
        // if (edges.depCnt > 0)
        //  detail::signal_edge(*this);
        edges = other.edges;
        other.edges.depCnt = 0;
      }
      JobContainer& operator=(JobContainer&& other) {
        GAIA_ASSERT(core::addressof(other) != this);
        cnt::ilist_item::operator=(GAIA_MOV(other));
        state = other.state.load();
        prio = other.prio;
        flags = other.flags;
        func = GAIA_MOV(other.func);
 
        // if (edges.depCnt > 0)
        //  detail::signal_edge(*this);
        edges = other.edges;
        other.edges.depCnt = 0;
 
        return *this;
      }
 
      GAIA_NODISCARD static JobContainer create(uint32_t index, uint32_t generation, void* pCtx) {
        auto* ctx = (JobAllocCtx*)pCtx;
 
        JobContainer jc{};
        jc.idx = index;
        jc.data.gen = generation;
        jc.prio = ctx->priority;
 
        return jc;
      }
 
      GAIA_NODISCARD static JobHandle handle(const JobContainer& jc) {
        return JobHandle(jc.idx, jc.data.gen, (jc.prio == JobPriority::Low) != 0);
      }
    };
 
    class JobManager {
      cnt::ilist<JobContainer, JobHandle> m_jobData;
 
    public:
      JobContainer& data(JobHandle jobHandle) {
        return m_jobData[jobHandle.id()];
      }
      const JobContainer& data(JobHandle jobHandle) const {
        return m_jobData[jobHandle.id()];
      }
 
      GAIA_NODISCARD JobHandle alloc_job(const Job& job) {
        JobAllocCtx ctx{job.priority};
 
        auto handle = m_jobData.alloc(&ctx);
        auto& j = m_jobData[handle.id()];
 
        // Make sure there is not state yet
        GAIA_ASSERT(j.state == 0 || j.state == JobState::Released);
 
        j.edges = {};
        j.prio = ctx.priority;
        j.state.store(0);
        j.func = job.func;
        j.flags = job.flags;
        return handle;
      }
 
      void free_job(JobHandle jobHandle) {
        auto& jobData = m_jobData.free(jobHandle);
        GAIA_ASSERT(done(jobData));
        jobData.state.store(JobState::Released);
      }
 
      void reset() {
        m_jobData.clear();
      }
 
      static void run(JobContainer& jobData) {
        if (jobData.func.operator bool())
          jobData.func();
 
        finalize(jobData);
      }
 
      static bool signal_edge(JobContainer& jobData) {
        // Subtract from dependency counter
        const auto state = jobData.state.fetch_sub(1) - 1;
 
        // If the job is not submitted, we can't accept it
        const auto s = state & JobState::STATE_BITS_MASK;
        if (s != JobState::Submitted)
          return false;
 
        // If the job still has some dependencies left we can't accept it
        const auto deps = state & JobState::DEP_BITS_MASK;
        return deps == 0;
      }
 
      static void free_edges(JobContainer& jobData) {
        // We only allocate an array for 2 and more dependencies
        if (jobData.edges.depCnt <= 1)
          return;
 
        mem::AllocHelper::free(jobData.edges.pDeps);
        // jobData.edges.depCnt = 0;
        // jobData.edges.pDeps = nullptr;
      }
 
      void dep(JobHandle jobFirst, JobHandle jobSecond) {
        dep(std::span(&jobFirst, 1), jobSecond);
      }
 
      void dep(std::span<JobHandle> jobsFirst, JobHandle jobSecond) {
        GAIA_ASSERT(!jobsFirst.empty());
 
        GAIA_PROF_SCOPE(JobManager::dep);
 
        auto& secondData = data(jobSecond);
 
#if GAIA_ASSERT_ENABLED
        GAIA_ASSERT(!busy(const_cast<const JobContainer&>(secondData)));
        for (auto jobFirst: jobsFirst) {
          const auto& firstData = data(jobFirst);
          GAIA_ASSERT(!busy(firstData));
        }
#endif
 
        for (auto jobFirst: jobsFirst)
          dep_internal(jobFirst, jobSecond);
 
        // Tell jobSecond that it has new dependencies
        // secondData.canWait = true;
        const uint32_t cnt = (uint32_t)jobsFirst.size();
        [[maybe_unused]] const uint32_t statePrev = secondData.state.fetch_add(cnt);
        GAIA_ASSERT((statePrev & JobState::DEP_BITS_MASK) < DEP_BITS_MASK - 1);
      }
 
      void dep_refresh(std::span<JobHandle> jobsFirst, JobHandle jobSecond) {
        GAIA_ASSERT(!jobsFirst.empty());
 
        GAIA_PROF_SCOPE(JobManager::dep_refresh);
 
        auto& secondData = data(jobSecond);
 
#if GAIA_ASSERT_ENABLED
        GAIA_ASSERT(!busy(const_cast<const JobContainer&>(secondData)));
        for (auto jobFirst: jobsFirst) {
          const auto& firstData = data(jobFirst);
          GAIA_ASSERT(!busy(firstData));
        }
 
        for (auto jobFirst: jobsFirst)
          dep_refresh_internal(jobFirst, jobSecond);
#endif
 
        // Tell jobSecond that it has new dependencies
        // secondData.canWait = true;
        const uint32_t cnt = (uint32_t)jobsFirst.size();
        [[maybe_unused]] const uint32_t statePrev = secondData.state.fetch_add(cnt);
        GAIA_ASSERT((statePrev & JobState::DEP_BITS_MASK) < DEP_BITS_MASK - 1);
      }
 
      static uint32_t submit(JobContainer& jobData) {
        [[maybe_unused]] const auto state = jobData.state.load() & JobState::STATE_BITS_MASK;
        GAIA_ASSERT(state < JobState::Submitted);
        const auto val = jobData.state.fetch_add(JobState::Submitted) + (uint32_t)JobState::Submitted;
#if GAIA_LOG_JOB_STATES
        GAIA_LOG_N("JobHandle %u.%u - SUBMITTED", jobData.idx, jobData.gen);
#endif
        return val;
      }
 
      static void processing(JobContainer& jobData) {
        GAIA_ASSERT(submitted(const_cast<const JobContainer&>(jobData)));
        jobData.state.store(JobState::Processing);
#if GAIA_LOG_JOB_STATES
        GAIA_LOG_N("JobHandle %u.%u - PROCESSING", jobData.idx, jobData.gen);
#endif
      }
 
      static void executing(JobContainer& jobData, uint32_t workerIdx) {
        GAIA_ASSERT(processing(const_cast<const JobContainer&>(jobData)));
        jobData.state.store(JobState::Executing | workerIdx);
#if GAIA_LOG_JOB_STATES
        GAIA_LOG_N("JobHandle %u.%u - EXECUTING", jobData.idx, jobData.gen);
#endif
      }
 
      static void finalize(JobContainer& jobData) {
        jobData.state.store(JobState::Done);
#if GAIA_LOG_JOB_STATES
        GAIA_LOG_N("JobHandle %u.%u - DONE", jobData.idx, jobData.gen);
#endif
      }
 
      static void reset_state(JobContainer& jobData) {
        [[maybe_unused]] const auto state = jobData.state.load() & JobState::STATE_BITS_MASK;
        // The job needs to be either clear or finalize for us to allow a reset
        GAIA_ASSERT(state == 0 || state == JobState::Done);
        jobData.state.store(0);
#if GAIA_LOG_JOB_STATES
        GAIA_LOG_N("JobHandle %u.%u - RESET_STATE", jobData.idx, jobData.gen);
#endif
      }
 
      GAIA_NODISCARD bool is_clear(JobHandle jobHandle) const {
        const auto& jobData = data(jobHandle);
        const auto state = jobData.state.load();
        return state == 0;
      }
 
      GAIA_NODISCARD static bool is_clear(JobContainer& jobData) {
        const auto state = jobData.state.load();
        return state == 0;
      }
 
      GAIA_NODISCARD static bool submitted(const JobContainer& jobData) {
        const auto state = jobData.state.load() & JobState::STATE_BITS_MASK;
        return state == JobState::Submitted;
      }
 
      GAIA_NODISCARD static bool processing(const JobContainer& jobData) {
        const auto state = jobData.state.load() & JobState::STATE_BITS_MASK;
        return state == JobState::Processing;
      }
 
      GAIA_NODISCARD static bool busy(const JobContainer& jobData) {
        const auto state = jobData.state.load() & JobState::STATE_BITS_MASK;
        return state == JobState::Executing || state == JobState::Processing;
      }
 
      GAIA_NODISCARD static bool done(const JobContainer& jobData) {
        const auto state = jobData.state.load() & JobState::STATE_BITS_MASK;
        return state == JobState::Done;
      }
 
    private:
      void dep_internal(JobHandle jobFirst, JobHandle jobSecond) {
        GAIA_ASSERT(jobFirst != (JobHandle)JobNull_t{});
        GAIA_ASSERT(jobSecond != (JobHandle)JobNull_t{});
 
        auto& firstData = data(jobFirst);
        const auto depCnt0 = firstData.edges.depCnt;
        const auto depCnt1 = ++firstData.edges.depCnt;
 
#if GAIA_LOG_JOB_STATES
        GAIA_LOG_N(
            "DEP %u.%u, %u -> %u.%u", jobFirst.id(), jobFirst.gen(), firstData.edges.depCnt, jobSecond.id(),
            jobSecond.gen());
#endif
 
        if (depCnt1 <= 1) {
          firstData.edges.dep = jobSecond;
        } else if (depCnt1 == 2) {
          auto prev = firstData.edges.dep;
          // TODO: Use custom allocator
          firstData.edges.pDeps = mem::AllocHelper::alloc<JobHandle>(depCnt1);
          firstData.edges.pDeps[0] = prev;
          firstData.edges.pDeps[1] = jobSecond;
        } else {
          // Reallocate if the previous value was a power of 2
          const bool isPow2 = core::is_pow2(depCnt0);
          if (isPow2) {
            const auto nextPow2 = depCnt0 << 1;
            auto* pPrev = firstData.edges.pDeps;
            // TODO: Use custom allocator
            firstData.edges.pDeps = mem::AllocHelper::alloc<JobHandle>(nextPow2);
            if (pPrev != nullptr) {
              GAIA_FOR(depCnt0) firstData.edges.pDeps[i] = pPrev[i];
              mem::AllocHelper::free(pPrev);
            }
          }
 
          // Append new dependencies
          firstData.edges.pDeps[depCnt0] = jobSecond;
        }
      }
 
#if GAIA_ASSERT_ENABLED
      void dep_refresh_internal(JobHandle jobFirst, JobHandle jobSecond) const {
        GAIA_ASSERT(jobFirst != (JobHandle)JobNull_t{});
        GAIA_ASSERT(jobSecond != (JobHandle)JobNull_t{});
 
        const auto& firstData = data(jobFirst);
        const auto depCnt = firstData.edges.depCnt;
 
        if (depCnt <= 1) {
          GAIA_ASSERT(firstData.edges.dep == jobSecond);
        } else {
          GAIA_ASSERT(firstData.edges.pDeps != nullptr);
          bool found = false;
          GAIA_FOR(firstData.edges.depCnt) {
            if (firstData.edges.pDeps[i] == jobSecond) {
              found = true;
              break;
            }
          }
          GAIA_ASSERT(found);
        }
      }
#endif
    };
 
    namespace detail {
      void signal_edge(JobContainer& jobData) {
        JobManager::signal_edge(jobData);
      }
    } // namespace detail
  } // namespace mt
} // namespace gaia