#include "wiJobSystem.h"

#include "wiSpinLock.h"

#include "wiBacklog.h"

#include "wiPlatform.h"

#include "wiTimer.h"

#include "wiAllocator.h"

#include <memory>

#include <algorithm>

#include <string>

#include <thread>

#include <mutex>

#include <condition_variable>

#ifdef _WIN32

#include <malloc.h> // alloca

#endif // _WIN32

#ifdef PLATFORM_LINUX

#include <pthread.h>

#ifdef __FREEBSD__

#include <pthread_np.h>

#endif

#include <sys/resource.h>

#endif // PLATFORM_LINUX

#ifdef PLATFORM_PS5

#include "wiJobSystem_PS5.h"

#endif // PLATFORM_PS5

#ifdef __APPLE__

#include <sys/qos.h>

#endif // __APPLE__

namespace wi::jobsystem

{

struct alignas(64) Job

{

job_function_type task;

context* ctx;

uint32_t groupID;

uint32_t groupJobOffset;

uint32_t groupJobEnd;

uint32_t sharedmemory_size;

inline uint32_t execute()

{

JobArgs args;

args.groupID = groupID;

if (sharedmemory_size > 0)

{

static constexpr uint32_t alignment = 64; // avx-512 alignment is assumed at max

args.sharedmemory = alloca(sharedmemory_size + alignment); // overestimated alignment to not overwrite after allocation from the aligned pointer

args.sharedmemory = (void*)align((uint64_t)args.sharedmemory, (uint64_t)alignment);

}

else

{

args.sharedmemory = nullptr;

}

for (uint32_t j = groupJobOffset; j < groupJobEnd; ++j)

{

args.jobIndex = j;

args.groupIndex = j - groupJobOffset;

args.isFirstJobInGroup = (j == groupJobOffset);

args.isLastJobInGroup = (j == groupJobEnd - 1);

task(args);

}

return ctx->counter.fetch_sub(1, std::memory_order_relaxed); // returns context counter's previous value

}

};

struct JobQueue

{

struct Block

{

uint32_t first_item = 0;

uint32_t last_item = 0;

Block* next = nullptr;

Job items[256];

};

wi::allocator::BlockAllocator<Block, 16> allocator;

Block* first_block = nullptr;

Block* last_block = nullptr;

std::mutex locker;

//wi::SpinLock locker;

std::atomic_uint32_t cnt{ 0 }; // for early exit, reduce contention on locker in job stealing scenario

JobQueue()

{

first_block = last_block = allocator.allocate();

}

__forceinline void push_back(const Job& item)

{

std::scoped_lock lock(locker);

cnt.fetch_add(1, std::memory_order_relaxed);

if (last_block->last_item == arraysize(Block::items))

{

// We ran out of items in the last block, so we need to allocate a new one:

last_block->next = allocator.allocate();

last_block = last_block->next;

}

last_block->items[last_block->last_item++] = item;

}

__forceinline bool pop_front(Job& item)

{

if (cnt.load(std::memory_order_relaxed) == 0)

return false;

std::scoped_lock lock(locker);

if (first_block->first_item == first_block->last_item)

{

// Here it means that the container is empty

return false;

}

item = std::move(first_block->items[first_block->first_item++]);

if (first_block->first_item == arraysize(Block::items))

{

// When we are here it means that the block was emptied

Block* next = first_block->next;

if (next == nullptr)

{

// No next block means there is only one block and it became empty after popping

// -> we can reset just the block

first_block->first_item = 0;

first_block->last_item = 0;

}

else

{

// There is a next block, we have to move to it

// -> the current block can be freed and reused

allocator.free(first_block);

first_block = next;

}

}

cnt.fetch_sub(1, std::memory_order_relaxed);

return true;

}

};

struct PriorityResources

{

uint32_t numThreads = 0;

wi::vector<std::thread> threads;

std::unique_ptr<JobQueue[]> jobQueuePerThread;

std::atomic<uint8_t> nextQueue{ 0 };

std::condition_variable sleepingCondition; // for workers that are sleeping

std::mutex sleepingMutex; // for workers that are sleeping

std::condition_variable waitingCondition; // for unblocking a Wait()

std::mutex waitingMutex; // for unblocking a Wait()

uint8_t mod_lut[256] = {}; // lookup table from atomic uint8_t -> threadID (avoiding modulo)

constexpr uint8_t constrain_queue_index(uint8_t idx) const

{

//idx = idx % numThreads;

idx = mod_lut[idx]; // this has the modulo precomputed at Initialize()

return idx;

}

inline uint8_t next_queue_index()

{

uint8_t idx = nextQueue.fetch_add(1, std::memory_order_relaxed);

return constrain_queue_index(idx);

}

inline JobQueue& next_queue()

{

return jobQueuePerThread[next_queue_index()];

}

// Start working on a job queue

// After the job queue is finished, it can switch to an other queue and steal jobs from there

inline void work(uint32_t startingQueue)

{

Job job;

for (uint32_t i = 0; i < numThreads; ++i)

{

JobQueue& job_queue = jobQueuePerThread[constrain_queue_index(startingQueue)];

while (job_queue.pop_front(job))

{

uint32_t progress_before = job.execute();

if (progress_before == 1)

{

// This is likely the last job because the counter was 1 before it was decremented in execute()

// So wake up the waiting threads here

std::unique_lock<std::mutex> lock(waitingMutex);

waitingCondition.notify_all();

}

}

startingQueue++; // go to next queue

}

}

};

// This structure is responsible to stop worker thread loops.

// Once this is destroyed, worker threads will be woken up and end their loops.

struct InternalState

{

uint32_t numCores = 0;

PriorityResources resources[int(Priority::Count)];

std::atomic_bool alive{ true };

void ShutDown()

{

if (IsShuttingDown())

return;

alive.store(false); // indicate that new jobs cannot be started from this point

bool wake_loop = true;

std::thread waker([&] {

while (wake_loop)

{

for (auto& x : resources)

{

x.sleepingCondition.notify_all(); // wakes up sleeping worker threads

}

}

});

for (auto& x : resources)

{

for (auto& thread : x.threads)

{

thread.join();

}

}

wake_loop = false;

waker.join();

for (auto& x : resources)

{

x.jobQueuePerThread.reset();

x.threads.clear();

x.numThreads = 0;

}

numCores = 0;

}

~InternalState()

{

ShutDown();

}

} static internal_state;

void Initialize(uint32_t maxThreadCount)

{

if (internal_state.numCores > 0)

return;

maxThreadCount = clamp(maxThreadCount, 1u, (uint32_t)arraysize(PriorityResources::mod_lut));

wi::Timer timer;

// Retrieve the number of hardware threads in this system:

internal_state.numCores = std::thread::hardware_concurrency();

for (int prio = 0; prio < int(Priority::Count); ++prio)

{

const Priority priority = (Priority)prio;

PriorityResources& res = internal_state.resources[prio];

// Calculate the actual number of worker threads we want:

switch (priority)

{

case Priority::High:

res.numThreads = internal_state.numCores - 1; // -1 for main thread

break;

case Priority::Low:

res.numThreads = internal_state.numCores - 2; // -1 for main thread, -1 for streaming

break;

case Priority::Streaming:

res.numThreads = 1;

break;

default:

assert(0);

break;

}

res.numThreads = clamp(res.numThreads, 1u, maxThreadCount);

res.jobQueuePerThread.reset(new JobQueue[res.numThreads]);

res.threads.reserve(res.numThreads);

// Precompute lookup table of modulos to avoid divs at runtime:

for (uint32_t i = 0; i < arraysize(res.mod_lut); ++i)

{

res.mod_lut[i] = i % res.numThreads;

}

for (uint32_t threadID = 0; threadID < res.numThreads; ++threadID)

{

std::thread& worker = res.threads.emplace_back([threadID, priority, &res] {

#if defined(__FREEBSD__)

// TODO: FreeBSD's setpriority is incompatible with the expected Linux non-standard behavior

#elif defined(PLATFORM_LINUX)

// from the sched(2) manpage:

// In the current [Linux 2.6.23+] implementation, each unit of

// difference in the nice values of two processes results in a

// factor of 1.25 in the degree to which the scheduler favors

// the higher priority process.

//

// so 3 would mean that other (prio 0) threads are around twice as important

switch (priority) {

case Priority::Low:

if (setpriority(PRIO_PROCESS, 0, 3) != 0)

{

perror("setpriority");

}

break;

case Priority::Streaming:

if (setpriority(PRIO_PROCESS, 0, 2) != 0)

{

perror("setpriority");

}

break;

case Priority::High:

// nothing to do

break;

default:

assert(0);

}

#elif defined(__APPLE__)

switch (priority)

{

case Priority::High:

pthread_set_qos_class_self_np(QOS_CLASS_USER_INTERACTIVE, 0);

break;

case Priority::Low:

pthread_set_qos_class_self_np(QOS_CLASS_UTILITY, 0);

break;

case Priority::Streaming:

pthread_set_qos_class_self_np(QOS_CLASS_BACKGROUND, 0);

break;

default:

assert(0);

}

#endif // PLATFORM_LINUX

while (internal_state.alive.load(std::memory_order_relaxed))

{

res.work(threadID);

// finished with jobs, put to sleep

std::unique_lock<std::mutex> lock(res.sleepingMutex);

res.sleepingCondition.wait(lock);

}

});

auto handle = worker.native_handle();

int core = threadID + 1; // put threads on increasing cores starting from 2nd

if (priority == Priority::Streaming)

{

// Put streaming to last core:

core = internal_state.numCores - 1 - threadID;

}

#ifdef _WIN32

// Do Windows-specific thread setup:

// Put each thread on to dedicated core:

DWORD_PTR affinityMask = 1ull << core;

DWORD_PTR affinity_result = SetThreadAffinityMask(handle, affinityMask);

assert(affinity_result > 0);

if (priority == Priority::High)

{

BOOL priority_result = SetThreadPriority(handle, THREAD_PRIORITY_NORMAL);

assert(priority_result != 0);

std::wstring wthreadname = L"wi::job_" + std::to_wstring(threadID);

HRESULT hr = SetThreadDescription(handle, wthreadname.c_str());

assert(SUCCEEDED(hr));

}

else if (priority == Priority::Low)

{

BOOL priority_result = SetThreadPriority(handle, THREAD_PRIORITY_LOWEST);

assert(priority_result != 0);

std::wstring wthreadname = L"wi::job_lo_" + std::to_wstring(threadID);

HRESULT hr = SetThreadDescription(handle, wthreadname.c_str());

assert(SUCCEEDED(hr));

}

else if (priority == Priority::Streaming)

{

BOOL priority_result = SetThreadPriority(handle, THREAD_PRIORITY_BELOW_NORMAL);

assert(priority_result != 0);

std::wstring wthreadname = L"wi::job_st_" + std::to_wstring(threadID);

HRESULT hr = SetThreadDescription(handle, wthreadname.c_str());

assert(SUCCEEDED(hr));

}

#elif defined(PLATFORM_LINUX)

#define handle_error_en(en, msg) \

do { errno = en; perror(msg); } while (0)

int ret;

cpu_set_t cpuset;

CPU_ZERO(&cpuset);

size_t cpusetsize = sizeof(cpuset);

CPU_SET(core, &cpuset);

ret = pthread_setaffinity_np(handle, cpusetsize, &cpuset);

if (ret != 0)

handle_error_en(ret, std::string(" pthread_setaffinity_np[" + std::to_string(threadID) + ']').c_str());

if (priority == Priority::High)

{

std::string thread_name = "wi::job_" + std::to_string(threadID);

ret = pthread_setname_np(handle, thread_name.c_str());

if (ret != 0)

handle_error_en(ret, std::string(" pthread_setname_np[" + std::to_string(threadID) + ']').c_str());

}

else if (priority == Priority::Low)

{

std::string thread_name = "wi::job_lo_" + std::to_string(threadID);

ret = pthread_setname_np(handle, thread_name.c_str());

if (ret != 0)

handle_error_en(ret, std::string(" pthread_setname_np[" + std::to_string(threadID) + ']').c_str());

// priority is set in the worker function

}

else if (priority == Priority::Streaming)

{

std::string thread_name = "wi::job_st_" + std::to_string(threadID);

ret = pthread_setname_np(handle, thread_name.c_str());

if (ret != 0)

handle_error_en(ret, std::string(" pthread_setname_np[" + std::to_string(threadID) + ']').c_str());

// priority is set in the worker function

}

#undef handle_error_en

#elif defined(PLATFORM_PS5)

wi::jobsystem::ps5::SetupWorker(worker, threadID, core, priority);

#endif // _WIN32

}

}

wilog("wi::jobsystem Initialized with %d cores in %.2f ms\n\tHigh priority threads: %d\n\tLow priority threads: %d\n\tStreaming threads: %d", internal_state.numCores, timer.elapsed(), GetThreadCount(Priority::High), GetThreadCount(Priority::Low), GetThreadCount(Priority::Streaming));

}

void ShutDown()

{

internal_state.ShutDown();

}

bool IsShuttingDown()

{

return internal_state.alive.load(std::memory_order_relaxed) == false;

}

uint32_t GetThreadCount(Priority priority)

{

return internal_state.resources[int(priority)].numThreads;

}

void Execute(context& ctx, const job_function_type& task)

{

PriorityResources& res = internal_state.resources[int(ctx.priority)];

// Context state is updated:

ctx.counter.fetch_add(1, std::memory_order_relaxed);

Job job;

job.ctx = &ctx;

job.task = task;

job.groupID = 0;

job.groupJobOffset = 0;

job.groupJobEnd = 1;

job.sharedmemory_size = 0;

if (res.numThreads < 1)

{

// If job system is not yet initialized, job will be executed immediately here instead of thread:

job.execute();

return;

}

res.next_queue().push_back(job);

res.sleepingCondition.notify_one();

}

void Dispatch(context& ctx, uint32_t jobCount, uint32_t groupSize, const job_function_type& task, size_t sharedmemory_size)

{

if (jobCount == 0 || groupSize == 0)

{

return;

}

PriorityResources& res = internal_state.resources[int(ctx.priority)];

const uint32_t groupCount = DispatchGroupCount(jobCount, groupSize);

// Context state is updated:

ctx.counter.fetch_add(groupCount, std::memory_order_relaxed);

Job job;

job.ctx = &ctx;

job.task = task;

job.sharedmemory_size = (uint32_t)sharedmemory_size;

for (uint32_t groupID = 0; groupID < groupCount; ++groupID)

{

// For each group, generate one real job:

job.groupID = groupID;

job.groupJobOffset = groupID * groupSize;

job.groupJobEnd = std::min(job.groupJobOffset + groupSize, jobCount);

if (res.numThreads < 1)

{

// If job system is not yet initialized, job will be executed immediately here instead of thread:

job.execute();

}

else

{

res.next_queue().push_back(job);

}

}

if (res.numThreads > 1)

{

res.sleepingCondition.notify_all();

}

}

uint32_t DispatchGroupCount(uint32_t jobCount, uint32_t groupSize)

{

// Calculate the amount of job groups to dispatch (overestimate, or "ceil"):

return (jobCount + groupSize - 1) / groupSize;

}

bool IsBusy(const context& ctx)

{

// Whenever the context label is greater than zero, it means that there is still work that needs to be done

return ctx.counter.load(std::memory_order_relaxed) > 0;

}

void Wait(const context& ctx)

{

if (IsBusy(ctx))

{

PriorityResources& res = internal_state.resources[int(ctx.priority)];

// Wake any threads that might be sleeping:

res.sleepingCondition.notify_all();

// work() will pick up any jobs that are on standby and execute them on this thread:

res.work(res.next_queue_index());

while (IsBusy(ctx))

{

// If we are here, then there are still remaining jobs that work() couldn't pick up.

// The thread enters a sleep until the !IsBusy() waitCondition is signaled

std::unique_lock<std::mutex> lock(res.waitingMutex);

if (IsBusy(ctx)) // check after locking, to not enter wait when it was completed after lock

{

res.waitingCondition.wait(lock, [&ctx] { return !IsBusy(ctx); });

}

}

}

}

uint32_t GetRemainingJobCount(const context& ctx)

{

return ctx.counter.load(std::memory_order_relaxed);

}

}