#pragma once
#include "mip_image.h"
#include <whiteout/interfaces.h>
#include <whiteout/utils/job_group.h>
#include <algorithm>
#include <functional>
#include <memory>
#include <vector>
namespace whiteout::textures::mipmap {
struct PipelineContext {
interfaces::WorkerPool* pool = nullptr;
interfaces::TimelineSemaphore* sem = nullptr;
interfaces::TimelineSemaphore::Value currentValue = 0;
};
using Stage = std::function<void(MipImage&)>;
using PoolStage = std::function<void(MipImage&, PipelineContext*)>;
using Filter = std::function<void(const MipImage&, MipImage&)>;
using PoolFilter = std::function<void(const MipImage&, MipImage&, PipelineContext*)>;
inline void submitSingleTask(PipelineContext* ctx, std::function<void()> fn) {
if (!ctx || !ctx->sem) {
fn();
return;
}
const auto waitVal = ctx->currentValue;
ctx->currentValue = ctx->sem->next();
interfaces::WorkerTask task;
task.fn = std::move(fn);
task.waitSemaphore = ctx->sem;
task.waitValue = waitVal;
task.signalSemaphore = ctx->sem;
task.signalValue = ctx->currentValue;
ctx->pool->submit(task);
}
constexpr u32 kMinRowsPerTile = 32;
template <typename RowFn>
void parallelForRows(u32 totalRows, PipelineContext* ctx, RowFn&& fn) {
if (!ctx || !ctx->pool) {
fn(0u, totalRows);
return;
}
const auto singleJobFn = [&]() {
if (!ctx->sem) {
fn(0u, totalRows);
return;
}
const auto waitVal = ctx->currentValue;
ctx->currentValue = ctx->sem->next();
const auto signalVal = ctx->currentValue;
interfaces::WorkerTask task;
task.fn = [totalRows, fn]() { fn(0u, totalRows); };
task.waitSemaphore = ctx->sem;
task.waitValue = waitVal;
task.signalSemaphore = ctx->sem;
task.signalValue = signalVal;
ctx->pool->submit(task);
};
if (totalRows <= kMinRowsPerTile) {
singleJobFn();
return;
}
const u32 threadCount = static_cast<u32>(ctx->pool->threadCount());
const u32 tilesWanted = std::min(threadCount * 2, totalRows); const u32 rowsPerTile = std::max(totalRows / tilesWanted, 1u);
const u32 tileCount = (totalRows + rowsPerTile - 1) / rowsPerTile;
if (tileCount <= 1) {
singleJobFn();
return;
}
if (ctx->sem) {
const auto waitVal = ctx->currentValue;
ctx->currentValue = ctx->sem->next();
const auto signalVal = ctx->currentValue;
auto jobGroup = std::make_shared<utils::JobGroup>();
jobGroup->add(tileCount);
jobGroup->signalOnComplete(ctx->sem, signalVal);
for (u32 t = 0; t < tileCount; ++t) {
const u32 startRow = t * rowsPerTile;
const u32 endRow = std::min(startRow + rowsPerTile, totalRows);
interfaces::WorkerTask task;
task.fn = [startRow, endRow, fn, jg = jobGroup]() {
fn(startRow, endRow);
jg->done();
};
task.waitSemaphore = ctx->sem;
task.waitValue = waitVal;
ctx->pool->submit(task);
}
} else {
utils::JobGroup jobGroup;
jobGroup.add(tileCount);
for (u32 t = 0; t < tileCount; ++t) {
const u32 startRow = t * rowsPerTile;
const u32 endRow = std::min(startRow + rowsPerTile, totalRows);
interfaces::WorkerTask task{[startRow, endRow, &fn, &jobGroup]() {
fn(startRow, endRow);
jobGroup.done();
}};
ctx->pool->submit(task);
}
jobGroup.wait();
}
}
struct MipmapPipeline {
std::vector<PoolStage>
preProcess; PoolFilter downsample; std::vector<PoolStage> postProcess;
MipImage execute(const MipImage& src, u32 targetWidth, u32 targetHeight,
PipelineContext* ctx = nullptr) const {
MipImage input = src;
for (const auto& stage : preProcess)
stage(input, ctx);
MipImage result(targetWidth, targetHeight, input.channels);
downsample(input, result, ctx);
for (const auto& stage : postProcess)
stage(result, ctx);
return result;
}
interfaces::TimelineSemaphore::Value executeAsync(
const MipImage& src, u32 targetWidth, u32 targetHeight, interfaces::WorkerPool* pool,
interfaces::TimelineSemaphore* sem, interfaces::TimelineSemaphore::Value startValue,
MipImage* output) const {
struct State {
MipImage input;
MipImage result;
};
auto state = std::make_shared<State>();
state->input = src;
state->result = MipImage(targetWidth, targetHeight, src.channels);
PipelineContext ctx;
ctx.pool = pool;
ctx.sem = sem;
ctx.currentValue = startValue;
for (const auto& stage : preProcess)
stage(state->input, &ctx);
downsample(state->input, state->result, &ctx);
for (const auto& stage : postProcess)
stage(state->result, &ctx);
const auto waitVal = ctx.currentValue;
const auto signalVal = sem->next();
ctx.currentValue = signalVal;
interfaces::WorkerTask copyTask;
copyTask.fn = [state, output]() { *output = std::move(state->result); };
copyTask.waitSemaphore = sem;
copyTask.waitValue = waitVal;
copyTask.signalSemaphore = sem;
copyTask.signalValue = signalVal;
pool->submit(copyTask);
return signalVal;
}
};
}