Files
dusklight/mods/ao_mod/src/mod.cpp
T
2026-07-10 01:10:23 -06:00

932 lines
35 KiB
C++

// Ambient occlusion (GTAO) example mod.
//
// Showcases the gfx service's compute tasks and the camera service: after opaque scene draws,
// before translucent/fog overlays, the scene depth is resolved and a three-dispatch compute
// chain (depth MIP prefilter, GTAO, spatial denoise) produces a visibility texture that a
// fullscreen draw multiplies over the world.
//
// The WGSL in res/ is ported from Bevy Engine's SSAO implementation (MIT OR Apache-2.0),
// itself based on Intel XeGTAO (MIT); see res/licenses/ and the `PORT:` notes in the shaders.
#include "mods/service.hpp"
#include "mods/svc/camera.h"
#include "mods/svc/config.h"
#include "mods/svc/gfx.h"
#include "mods/svc/log.h"
#include "mods/svc/resource.h"
#include "mods/svc/ui.h"
#include <algorithm>
#include <atomic>
#include <cstring>
#include <initializer_list>
#include <type_traits>
#include <utility>
#include <vector>
#include <webgpu/webgpu.h>
DEFINE_MOD();
IMPORT_SERVICE(LogService, svc_log);
IMPORT_SERVICE(ConfigService, svc_config);
IMPORT_SERVICE(ResourceService, svc_resource);
IMPORT_SERVICE(UiService, svc_ui);
IMPORT_SERVICE(GfxService, svc_gfx);
IMPORT_SERVICE(CameraService, svc_camera);
namespace {
ConfigVarHandle g_cvarEnabled = 0;
ConfigVarHandle g_cvarQuality = 0;
ConfigVarHandle g_cvarRadius = 0;
ConfigVarHandle g_cvarIntensity = 0;
ConfigVarHandle g_cvarHalfRes = 0;
ConfigVarHandle g_cvarDebugView = 0;
GfxComputeTypeHandle g_computeType = 0;
GfxDrawTypeHandle g_drawType = 0;
GfxStageHookHandle g_afterOpaqueHook = 0;
UiWindowHandle g_controlsWindow = 0;
ResourceBuffer g_preprocessSource = RESOURCE_BUFFER_INIT;
ResourceBuffer g_gtaoSource = RESOURCE_BUFFER_INIT;
ResourceBuffer g_denoiseSource = RESOURCE_BUFFER_INIT;
ResourceBuffer g_compositeSource = RESOURCE_BUFFER_INIT;
GfxDeviceInfo g_deviceInfo = GFX_DEVICE_INFO_INIT;
WGPUComputePipeline g_preprocessPipeline = nullptr;
WGPUComputePipeline g_mip4Pipeline = nullptr;
WGPUComputePipeline g_gtaoPipeline = nullptr;
WGPUComputePipeline g_denoisePipeline = nullptr;
WGPUBindGroupLayout g_preprocessLayout = nullptr;
WGPUBindGroupLayout g_mip4Layout = nullptr;
WGPUBindGroupLayout g_gtaoLayout = nullptr;
WGPUBindGroupLayout g_denoiseLayout = nullptr;
WGPURenderPipeline g_compositePipeline = nullptr;
WGPURenderPipeline g_compositeDebugPipeline = nullptr;
WGPUBindGroupLayout g_compositeLayout = nullptr;
WGPUBindGroupLayout g_compositeDebugLayout = nullptr;
WGPUTexture g_hilbertLut = nullptr;
WGPUTextureView g_hilbertLutView = nullptr;
// AO chain targets, recreated when the render size (or halfRes) changes. Old sets are retired
// for a few frames instead of released immediately: payloads embedding their views may still
// be in flight on the render worker.
struct AoTargets {
uint32_t width = 0;
uint32_t height = 0;
WGPUTexture preprocessedDepth = nullptr;
WGPUTextureView preprocessedDepthMips[5] = {};
WGPUTextureView preprocessedDepthAll = nullptr;
WGPUTexture aoNoisy = nullptr;
WGPUTextureView aoNoisyView = nullptr;
WGPUTexture depthDifferences = nullptr;
WGPUTextureView depthDifferencesView = nullptr;
WGPUTexture aoFinal = nullptr;
WGPUTextureView aoFinalView = nullptr;
};
AoTargets g_targets;
struct RetiredTargets {
AoTargets targets;
int framesLeft = 0;
};
std::vector<RetiredTargets> g_retiredTargets;
bool g_warnedNoDepth = false;
bool g_loggedChain = false;
std::atomic g_chainExecuted{false};
// Mirror of the WGSL Uniforms struct (keep in sync with res/*.wgsl).
struct AoUniforms {
float projection[16];
float inverse_projection[16];
float size[2];
float inv_size[2];
float depth_scale[2];
float effect_radius;
float intensity;
float slice_count;
float samples_per_slice_side;
uint32_t debug_view;
float _pad;
};
static_assert(sizeof(AoUniforms) % 16 == 0);
struct ComputePayload {
WGPUTextureView depth; // frame-pooled scene depth snapshot
WGPUTextureView preprocessedDepthMips[5];
WGPUTextureView preprocessedDepthAll;
WGPUTextureView aoNoisy;
WGPUTextureView depthDifferences;
WGPUTextureView aoFinal;
uint32_t uniform_offset;
uint32_t uniform_size;
uint32_t width;
uint32_t height;
};
static_assert(sizeof(ComputePayload) <= GFX_INLINE_DRAW_PAYLOAD_SIZE);
static_assert(std::is_trivially_copyable_v<ComputePayload>);
struct CompositePayload {
WGPUTextureView aoFinal;
WGPUTextureView preprocessedDepth; // debug views reconstruct normals/depth from it
WGPUTextureView sceneDepth; // raw snapshot, for the bypass debug views
uint32_t uniform_offset;
uint32_t uniform_size;
uint32_t debug_view;
};
static_assert(sizeof(CompositePayload) <= GFX_INLINE_DRAW_PAYLOAD_SIZE);
static_assert(std::is_trivially_copyable_v<CompositePayload>);
int64_t get_int_option(ConfigVarHandle handle, int64_t fallback) {
int64_t value = fallback;
if (handle == 0 || svc_config->get_int(mod_ctx, handle, &value) != MOD_OK) {
return fallback;
}
return value;
}
bool get_bool_option(ConfigVarHandle handle, bool fallback) {
bool value = fallback;
if (handle == 0 || svc_config->get_bool(mod_ctx, handle, &value) != MOD_OK) {
return fallback;
}
return value;
}
// XeGTAO/Bevy quality presets: slices x (samples per slice side * 2).
void quality_counts(int64_t quality, float& sliceCount, float& samplesPerSliceSide) {
switch (std::clamp<int64_t>(quality, 0, 3)) {
case 0:
sliceCount = 1.0f;
samplesPerSliceSide = 2.0f;
break;
case 1:
sliceCount = 2.0f;
samplesPerSliceSide = 2.0f;
break;
default:
case 2:
sliceCount = 3.0f;
samplesPerSliceSide = 3.0f;
break;
case 3:
sliceCount = 9.0f;
samplesPerSliceSide = 3.0f;
break;
}
}
WGPUShaderModule create_shader_module(const char* label, const ResourceBuffer& source) {
WGPUShaderSourceWGSL wgsl = WGPU_SHADER_SOURCE_WGSL_INIT;
wgsl.code = {static_cast<const char*>(source.data), source.size};
WGPUShaderModuleDescriptor moduleDesc = WGPU_SHADER_MODULE_DESCRIPTOR_INIT;
moduleDesc.nextInChain = &wgsl.chain;
moduleDesc.label = {label, WGPU_STRLEN};
return wgpuDeviceCreateShaderModule(g_deviceInfo.device, &moduleDesc);
}
bool build_compute_pipeline(const char* label, const ResourceBuffer& source, const char* entry,
WGPUComputePipeline& outPipeline, WGPUBindGroupLayout& outLayout) {
WGPUShaderModule module = create_shader_module(label, source);
if (module == nullptr) {
return false;
}
WGPUComputePipelineDescriptor pipelineDesc = WGPU_COMPUTE_PIPELINE_DESCRIPTOR_INIT;
pipelineDesc.label = {label, WGPU_STRLEN};
pipelineDesc.compute.module = module;
pipelineDesc.compute.entryPoint = {entry, WGPU_STRLEN};
outPipeline = wgpuDeviceCreateComputePipeline(g_deviceInfo.device, &pipelineDesc);
wgpuShaderModuleRelease(module);
if (outPipeline == nullptr) {
return false;
}
outLayout = wgpuComputePipelineGetBindGroupLayout(outPipeline, 0);
return outLayout != nullptr;
}
bool build_composite_pipeline(
bool blend, WGPURenderPipeline& outPipeline, WGPUBindGroupLayout& outLayout) {
WGPUShaderModule module = create_shader_module("AO composite", g_compositeSource);
if (module == nullptr) {
return false;
}
// Multiply blend
WGPUBlendState blendState{
.color =
{
.operation = WGPUBlendOperation_Add,
.srcFactor = WGPUBlendFactor_Dst,
.dstFactor = WGPUBlendFactor_Zero,
},
.alpha =
{
.operation = WGPUBlendOperation_Add,
.srcFactor = WGPUBlendFactor_Zero,
.dstFactor = WGPUBlendFactor_One,
},
};
WGPUColorTargetState colorTarget = WGPU_COLOR_TARGET_STATE_INIT;
colorTarget.format = g_deviceInfo.color_format;
if (blend) {
colorTarget.blend = &blendState;
}
WGPUFragmentState fragment = WGPU_FRAGMENT_STATE_INIT;
fragment.module = module;
fragment.entryPoint = {"fs_main", WGPU_STRLEN};
fragment.targetCount = 1;
fragment.targets = &colorTarget;
// Depth state must match the EFB pass despite never touching depth.
WGPUDepthStencilState depthStencil = WGPU_DEPTH_STENCIL_STATE_INIT;
depthStencil.format = g_deviceInfo.depth_format;
depthStencil.depthWriteEnabled = WGPUOptionalBool_False;
depthStencil.depthCompare = WGPUCompareFunction_Always;
WGPURenderPipelineDescriptor pipelineDesc = WGPU_RENDER_PIPELINE_DESCRIPTOR_INIT;
pipelineDesc.label = {blend ? "AO composite" : "AO composite (debug)", WGPU_STRLEN};
pipelineDesc.vertex.module = module;
pipelineDesc.vertex.entryPoint = {"vs_main", WGPU_STRLEN};
pipelineDesc.primitive.topology = WGPUPrimitiveTopology_TriangleList;
pipelineDesc.depthStencil = &depthStencil;
pipelineDesc.multisample.count = g_deviceInfo.sample_count;
pipelineDesc.fragment = &fragment;
outPipeline = wgpuDeviceCreateRenderPipeline(g_deviceInfo.device, &pipelineDesc);
wgpuShaderModuleRelease(module);
if (outPipeline == nullptr) {
return false;
}
outLayout = wgpuRenderPipelineGetBindGroupLayout(outPipeline, 0);
return outLayout != nullptr;
}
// Hilbert curve index LUT for the R2 noise sequence, generated once at init.
// Ported from Bevy's generate_hilbert_index_lut (https://www.shadertoy.com/view/3tB3z3).
uint16_t hilbert_index(uint16_t x, uint16_t y) {
uint16_t index = 0;
for (uint16_t level = 32; level > 0; level /= 2) {
const uint16_t regionX = (x & level) > 0 ? 1 : 0;
const uint16_t regionY = (y & level) > 0 ? 1 : 0;
index += level * level * ((3 * regionX) ^ regionY);
if (regionY == 0) {
if (regionX == 1) {
x = 63 - x;
y = 63 - y;
}
std::swap(x, y);
}
}
return index;
}
bool build_hilbert_lut() {
WGPUTextureDescriptor texDesc = WGPU_TEXTURE_DESCRIPTOR_INIT;
texDesc.label = {"AO hilbert LUT", WGPU_STRLEN};
texDesc.usage = WGPUTextureUsage_TextureBinding | WGPUTextureUsage_CopyDst;
texDesc.size = {64, 64, 1};
texDesc.format = WGPUTextureFormat_R16Uint;
g_hilbertLut = wgpuDeviceCreateTexture(g_deviceInfo.device, &texDesc);
if (g_hilbertLut == nullptr) {
return false;
}
g_hilbertLutView = wgpuTextureCreateView(g_hilbertLut, nullptr);
if (g_hilbertLutView == nullptr) {
return false;
}
uint16_t lut[64 * 64];
for (uint16_t y = 0; y < 64; ++y) {
for (uint16_t x = 0; x < 64; ++x) {
lut[y * 64 + x] = hilbert_index(x, y);
}
}
WGPUTexelCopyTextureInfo dst = WGPU_TEXEL_COPY_TEXTURE_INFO_INIT;
dst.texture = g_hilbertLut;
WGPUTexelCopyBufferLayout layout{.offset = 0, .bytesPerRow = 64 * 2, .rowsPerImage = 64};
WGPUExtent3D extent{64, 64, 1};
wgpuQueueWriteTexture(g_deviceInfo.queue, &dst, lut, sizeof(lut), &layout, &extent);
return true;
}
void release_targets(AoTargets& targets) {
for (auto*& view : targets.preprocessedDepthMips) {
if (view != nullptr) {
wgpuTextureViewRelease(view);
view = nullptr;
}
}
const auto releaseView = [](WGPUTextureView& view) {
if (view != nullptr) {
wgpuTextureViewRelease(view);
view = nullptr;
}
};
const auto releaseTexture = [](WGPUTexture& texture) {
if (texture != nullptr) {
wgpuTextureRelease(texture);
texture = nullptr;
}
};
releaseView(targets.preprocessedDepthAll);
releaseView(targets.aoNoisyView);
releaseView(targets.depthDifferencesView);
releaseView(targets.aoFinalView);
releaseTexture(targets.preprocessedDepth);
releaseTexture(targets.aoNoisy);
releaseTexture(targets.depthDifferences);
releaseTexture(targets.aoFinal);
targets.width = targets.height = 0;
}
void tick_retired_targets() {
for (auto it = g_retiredTargets.begin(); it != g_retiredTargets.end();) {
if (--it->framesLeft <= 0) {
release_targets(it->targets);
it = g_retiredTargets.erase(it);
} else {
++it;
}
}
}
bool ensure_targets(uint32_t width, uint32_t height) {
if (g_targets.width == width && g_targets.height == height) {
return true;
}
if (g_targets.width != 0) {
g_retiredTargets.push_back(RetiredTargets{std::exchange(g_targets, AoTargets{}), 4});
}
const auto createStorageTexture = [&](const char* label, WGPUTextureFormat format,
uint32_t mipCount, WGPUTexture& outTexture) {
WGPUTextureDescriptor texDesc = WGPU_TEXTURE_DESCRIPTOR_INIT;
texDesc.label = {label, WGPU_STRLEN};
texDesc.usage = WGPUTextureUsage_StorageBinding | WGPUTextureUsage_TextureBinding;
texDesc.size = {width, height, 1};
texDesc.format = format;
texDesc.mipLevelCount = mipCount;
outTexture = wgpuDeviceCreateTexture(g_deviceInfo.device, &texDesc);
return outTexture != nullptr;
};
bool ok = createStorageTexture("AO preprocessed depth", WGPUTextureFormat_R32Float, 5,
g_targets.preprocessedDepth) &&
createStorageTexture("AO noisy", WGPUTextureFormat_R32Float, 1, g_targets.aoNoisy) &&
createStorageTexture("AO depth differences", WGPUTextureFormat_R32Uint, 1,
g_targets.depthDifferences) &&
createStorageTexture("AO final", WGPUTextureFormat_R32Float, 1, g_targets.aoFinal);
if (ok) {
for (uint32_t mip = 0; mip < 5 && ok; ++mip) {
WGPUTextureViewDescriptor viewDesc = WGPU_TEXTURE_VIEW_DESCRIPTOR_INIT;
viewDesc.baseMipLevel = mip;
viewDesc.mipLevelCount = 1;
g_targets.preprocessedDepthMips[mip] =
wgpuTextureCreateView(g_targets.preprocessedDepth, &viewDesc);
ok = g_targets.preprocessedDepthMips[mip] != nullptr;
}
}
if (ok) {
g_targets.preprocessedDepthAll =
wgpuTextureCreateView(g_targets.preprocessedDepth, nullptr);
g_targets.aoNoisyView = wgpuTextureCreateView(g_targets.aoNoisy, nullptr);
g_targets.depthDifferencesView = wgpuTextureCreateView(g_targets.depthDifferences, nullptr);
g_targets.aoFinalView = wgpuTextureCreateView(g_targets.aoFinal, nullptr);
ok = g_targets.preprocessedDepthAll != nullptr && g_targets.aoNoisyView != nullptr &&
g_targets.depthDifferencesView != nullptr && g_targets.aoFinalView != nullptr;
}
if (!ok) {
release_targets(g_targets);
return false;
}
g_targets.width = width;
g_targets.height = height;
return true;
}
constexpr uint32_t div_ceil(uint32_t numerator, uint32_t denominator) {
return (numerator + denominator - 1) / denominator;
}
// Render worker thread: the AO chain as one compute pass with three dispatches.
void on_compute(
ModContext*, const GfxComputeContext* ctx, const void* payload, size_t payloadSize, void*) {
if (payloadSize != sizeof(ComputePayload)) {
return;
}
ComputePayload data;
std::memcpy(&data, payload, sizeof(data));
if (data.depth == nullptr || g_preprocessPipeline == nullptr) {
return;
}
const auto makeBindGroup = [&](WGPUBindGroupLayout layout,
std::initializer_list<WGPUBindGroupEntry> entries) {
WGPUBindGroupDescriptor bindGroupDesc = WGPU_BIND_GROUP_DESCRIPTOR_INIT;
bindGroupDesc.layout = layout;
bindGroupDesc.entryCount = entries.size();
bindGroupDesc.entries = entries.begin();
return wgpuDeviceCreateBindGroup(ctx->device, &bindGroupDesc);
};
const auto textureEntry = [](uint32_t binding, WGPUTextureView view) {
WGPUBindGroupEntry entry = WGPU_BIND_GROUP_ENTRY_INIT;
entry.binding = binding;
entry.textureView = view;
return entry;
};
const auto uniformEntry = [&](uint32_t binding) {
WGPUBindGroupEntry entry = WGPU_BIND_GROUP_ENTRY_INIT;
entry.binding = binding;
entry.buffer = ctx->uniform_buffer;
entry.offset = data.uniform_offset;
entry.size = data.uniform_size;
return entry;
};
WGPUBindGroup preprocessGroup = makeBindGroup(g_preprocessLayout,
{textureEntry(0, data.depth), textureEntry(1, data.preprocessedDepthMips[0]),
textureEntry(2, data.preprocessedDepthMips[1]),
textureEntry(3, data.preprocessedDepthMips[2]),
textureEntry(4, data.preprocessedDepthMips[3]), uniformEntry(5)});
WGPUBindGroup mip4Group =
makeBindGroup(g_mip4Layout, {textureEntry(6, data.preprocessedDepthMips[3]),
textureEntry(7, data.preprocessedDepthMips[4])});
WGPUBindGroup gtaoGroup = makeBindGroup(
g_gtaoLayout, {textureEntry(0, data.preprocessedDepthAll),
textureEntry(1, g_hilbertLutView), textureEntry(2, data.aoNoisy),
textureEntry(3, data.depthDifferences), uniformEntry(4)});
WGPUBindGroup denoiseGroup = makeBindGroup(
g_denoiseLayout, {textureEntry(0, data.aoNoisy), textureEntry(1, data.depthDifferences),
textureEntry(2, data.aoFinal), uniformEntry(3)});
if (preprocessGroup == nullptr || mip4Group == nullptr || gtaoGroup == nullptr ||
denoiseGroup == nullptr)
{
const auto release = [](WGPUBindGroup group) {
if (group != nullptr) {
wgpuBindGroupRelease(group);
}
};
release(preprocessGroup);
release(mip4Group);
release(gtaoGroup);
release(denoiseGroup);
return;
}
WGPUComputePassDescriptor passDesc = WGPU_COMPUTE_PASS_DESCRIPTOR_INIT;
passDesc.label = {"AO chain", WGPU_STRLEN};
WGPUComputePassEncoder pass = wgpuCommandEncoderBeginComputePass(ctx->encoder, &passDesc);
// Each preprocess workgroup covers 16x16 MIP-0 texels (8x8 invocations, 2x2 texels each).
wgpuComputePassEncoderSetPipeline(pass, g_preprocessPipeline);
wgpuComputePassEncoderSetBindGroup(pass, 0, preprocessGroup, 0, nullptr);
wgpuComputePassEncoderDispatchWorkgroups(
pass, div_ceil(data.width, 16), div_ceil(data.height, 16), 1);
wgpuComputePassEncoderSetPipeline(pass, g_mip4Pipeline);
wgpuComputePassEncoderSetBindGroup(pass, 0, mip4Group, 0, nullptr);
wgpuComputePassEncoderDispatchWorkgroups(pass, div_ceil(std::max(data.width >> 4, 1u), 8),
div_ceil(std::max(data.height >> 4, 1u), 8), 1);
wgpuComputePassEncoderSetPipeline(pass, g_gtaoPipeline);
wgpuComputePassEncoderSetBindGroup(pass, 0, gtaoGroup, 0, nullptr);
wgpuComputePassEncoderDispatchWorkgroups(
pass, div_ceil(data.width, 8), div_ceil(data.height, 8), 1);
wgpuComputePassEncoderSetPipeline(pass, g_denoisePipeline);
wgpuComputePassEncoderSetBindGroup(pass, 0, denoiseGroup, 0, nullptr);
wgpuComputePassEncoderDispatchWorkgroups(
pass, div_ceil(data.width, 8), div_ceil(data.height, 8), 1);
wgpuComputePassEncoderEnd(pass);
wgpuComputePassEncoderRelease(pass);
wgpuBindGroupRelease(preprocessGroup);
wgpuBindGroupRelease(mip4Group);
wgpuBindGroupRelease(gtaoGroup);
wgpuBindGroupRelease(denoiseGroup);
g_chainExecuted.store(true, std::memory_order_release);
}
// Render worker thread: composite the AO over the scene (or show it, in debug view).
void on_draw(
ModContext*, const GfxDrawContext* ctx, const void* payload, size_t payloadSize, void*) {
if (payloadSize != sizeof(CompositePayload)) {
return;
}
CompositePayload data;
std::memcpy(&data, payload, sizeof(data));
WGPURenderPipeline pipeline =
data.debug_view != 0 ? g_compositeDebugPipeline : g_compositePipeline;
WGPUBindGroupLayout layout = data.debug_view != 0 ? g_compositeDebugLayout : g_compositeLayout;
if (data.aoFinal == nullptr || data.preprocessedDepth == nullptr ||
data.sceneDepth == nullptr || pipeline == nullptr)
{
return;
}
WGPUBindGroupEntry entries[4] = {WGPU_BIND_GROUP_ENTRY_INIT, WGPU_BIND_GROUP_ENTRY_INIT,
WGPU_BIND_GROUP_ENTRY_INIT, WGPU_BIND_GROUP_ENTRY_INIT};
entries[0].binding = 0;
entries[0].textureView = data.aoFinal;
entries[1].binding = 1;
entries[1].textureView = data.preprocessedDepth;
entries[2].binding = 2;
entries[2].textureView = data.sceneDepth;
entries[3].binding = 3;
entries[3].buffer = ctx->uniform_buffer;
entries[3].offset = data.uniform_offset;
entries[3].size = data.uniform_size;
WGPUBindGroupDescriptor bindGroupDesc = WGPU_BIND_GROUP_DESCRIPTOR_INIT;
bindGroupDesc.layout = layout;
bindGroupDesc.entryCount = 4;
bindGroupDesc.entries = entries;
WGPUBindGroup bindGroup = wgpuDeviceCreateBindGroup(ctx->device, &bindGroupDesc);
if (bindGroup == nullptr) {
return;
}
wgpuRenderPassEncoderSetPipeline(ctx->pass, pipeline);
wgpuRenderPassEncoderSetBindGroup(ctx->pass, 0, bindGroup, 0, nullptr);
wgpuRenderPassEncoderDraw(ctx->pass, 3, 1, 0, 0);
wgpuBindGroupRelease(bindGroup);
}
// Game thread, after opaque scene draws and before translucent/fog overlay lists.
void on_scene_after_opaque(ModContext*, const GfxStageContext* stageCtx, void*) {
tick_retired_targets();
if (!get_bool_option(g_cvarEnabled, true)) {
return;
}
if (stageCtx == nullptr || stageCtx->struct_size < sizeof(GfxStageContext) ||
stageCtx->game_view == nullptr)
{
return;
}
CameraInfo camera = CAMERA_INFO_INIT;
if (svc_camera->get_camera(mod_ctx, stageCtx->game_view, &camera) != MOD_OK) {
return;
}
GfxResolveDesc resolveDesc = GFX_RESOLVE_DESC_INIT;
resolveDesc.color = false;
resolveDesc.depth = true;
GfxResolvedTargets resolved = GFX_RESOLVED_TARGETS_INIT;
if (svc_gfx->resolve_pass(mod_ctx, &resolveDesc, &resolved) != MOD_OK ||
resolved.depth == nullptr)
{
if (!g_warnedNoDepth) {
g_warnedNoDepth = true;
svc_log->warn(mod_ctx, "depth snapshots unavailable; AO disabled");
}
return;
}
const bool halfRes = get_bool_option(g_cvarHalfRes, false);
const uint32_t divisor = halfRes ? 2 : 1;
const uint32_t width = resolved.width / divisor;
const uint32_t height = resolved.height / divisor;
if (width < 32 || height < 32 || !ensure_targets(width, height)) {
return;
}
AoUniforms uniforms{};
std::memcpy(uniforms.projection, camera.proj_from_view, sizeof(uniforms.projection));
std::memcpy(
uniforms.inverse_projection, camera.view_from_proj, sizeof(uniforms.inverse_projection));
uniforms.size[0] = static_cast<float>(width);
uniforms.size[1] = static_cast<float>(height);
uniforms.inv_size[0] = 1.0f / uniforms.size[0];
uniforms.inv_size[1] = 1.0f / uniforms.size[1];
uniforms.depth_scale[0] = static_cast<float>(resolved.width) / uniforms.size[0];
uniforms.depth_scale[1] = static_cast<float>(resolved.height) / uniforms.size[1];
uniforms.effect_radius =
static_cast<float>(std::clamp<int64_t>(get_int_option(g_cvarRadius, 70), 10, 500));
uniforms.intensity =
static_cast<float>(std::clamp<int64_t>(get_int_option(g_cvarIntensity, 100), 0, 100)) /
100.0f;
quality_counts(
get_int_option(g_cvarQuality, 2), uniforms.slice_count, uniforms.samples_per_slice_side);
const uint32_t debugMode =
static_cast<uint32_t>(std::clamp<int64_t>(get_int_option(g_cvarDebugView, 0), 0, 4));
uniforms.debug_view = debugMode;
GfxRange uniformRange{0, 0};
if (svc_gfx->push_uniform(mod_ctx, &uniforms, sizeof(uniforms), &uniformRange) != MOD_OK) {
return;
}
ComputePayload computePayload{};
computePayload.depth = resolved.depth;
for (int mip = 0; mip < 5; ++mip) {
computePayload.preprocessedDepthMips[mip] = g_targets.preprocessedDepthMips[mip];
}
computePayload.preprocessedDepthAll = g_targets.preprocessedDepthAll;
computePayload.aoNoisy = g_targets.aoNoisyView;
computePayload.depthDifferences = g_targets.depthDifferencesView;
computePayload.aoFinal = g_targets.aoFinalView;
computePayload.uniform_offset = uniformRange.offset;
computePayload.uniform_size = uniformRange.size;
computePayload.width = width;
computePayload.height = height;
if (svc_gfx->push_compute(mod_ctx, g_computeType, &computePayload, sizeof(computePayload)) !=
MOD_OK)
{
return;
}
const CompositePayload drawPayload{g_targets.aoFinalView, g_targets.preprocessedDepthAll,
resolved.depth, uniformRange.offset, uniformRange.size, debugMode};
svc_gfx->push_draw(mod_ctx, g_drawType, &drawPayload, sizeof(drawPayload));
}
void add_control(UiElementHandle pane, const UiControlDesc& desc) {
svc_ui->pane_add_control(mod_ctx, pane, &desc, nullptr);
}
void add_toggle(UiElementHandle pane, const char* label, ConfigVarHandle cvar, const char* help) {
UiControlDesc control = UI_CONTROL_DESC_INIT;
control.kind = UI_CONTROL_TOGGLE;
control.label = label;
control.help_rml = help;
control.binding = UI_BINDING_CONFIG_VAR;
control.config_var = cvar;
add_control(pane, control);
}
ModResult build_controls_tab(
ModContext*, UiWindowHandle, UiElementHandle left, UiElementHandle right, void*, ModError*) {
(void)right;
svc_ui->pane_add_section(mod_ctx, left, "Ambient Occlusion");
add_toggle(left, "Enabled", g_cvarEnabled, "Enables the GTAO pass.");
static const char* kQualityOptions[] = {"Low", "Medium", "High", "Ultra"};
UiControlDesc control = UI_CONTROL_DESC_INIT;
control.kind = UI_CONTROL_SELECT;
control.label = "Quality";
control.help_rml = "Horizon slices and samples per pixel (XeGTAO presets: 4/8/18/54 spp).";
control.binding = UI_BINDING_CONFIG_VAR;
control.config_var = g_cvarQuality;
control.options = kQualityOptions;
control.option_count = 4;
add_control(left, control);
control = UI_CONTROL_DESC_INIT;
control.kind = UI_CONTROL_NUMBER;
control.label = "Radius";
control.help_rml = "Occlusion sampling radius in world units.";
control.binding = UI_BINDING_CONFIG_VAR;
control.config_var = g_cvarRadius;
control.min = 10;
control.max = 500;
control.step = 10;
add_control(left, control);
control = UI_CONTROL_DESC_INIT;
control.kind = UI_CONTROL_NUMBER;
control.label = "Intensity";
control.help_rml = "How strongly occlusion darkens the scene.";
control.binding = UI_BINDING_CONFIG_VAR;
control.config_var = g_cvarIntensity;
control.min = 0;
control.max = 100;
control.step = 5;
control.suffix = "%";
add_control(left, control);
add_toggle(left, "Half Resolution", g_cvarHalfRes,
"Computes AO at half resolution and upscales; faster, slightly softer.");
static const char* kDebugOptions[] = {"Off", "AO", "Normals", "Depth", "Staircase"};
control = UI_CONTROL_DESC_INIT;
control.kind = UI_CONTROL_SELECT;
control.label = "Debug View";
control.help_rml = "AO: raw visibility as grayscale.<br/>Normals: the view-space "
"normals the GTAO pass consumes.<br/>Depth: the preprocessed depth "
"as a distance gradient.<br/>Staircase: detects quantized depth - smooth "
"depth is near-black with thin triangle edges, quantized depth lights "
"up across surfaces.";
control.binding = UI_BINDING_CONFIG_VAR;
control.config_var = g_cvarDebugView;
control.options = kDebugOptions;
control.option_count = 5;
add_control(left, control);
return MOD_OK;
}
void on_controls_window_closed(ModContext*, UiWindowHandle, void*) {
g_controlsWindow = 0;
}
void on_open_controls(ModContext*, void*) {
if (g_controlsWindow != 0) {
return;
}
UiTabDesc tabs[1] = {UI_TAB_DESC_INIT};
tabs[0].title = "Controls";
tabs[0].build = build_controls_tab;
UiWindowDesc desc = UI_WINDOW_DESC_INIT;
desc.tabs = tabs;
desc.tab_count = 1;
desc.on_closed = on_controls_window_closed;
if (svc_ui->window_push(mod_ctx, &desc, &g_controlsWindow) != MOD_OK) {
svc_log->error(mod_ctx, "failed to open AO controls window");
}
}
ModResult build_panel(ModContext*, UiElementHandle panel, void*, ModError*) {
UiControlDesc control = UI_CONTROL_DESC_INIT;
control.kind = UI_CONTROL_TOGGLE;
control.label = "Enabled";
control.binding = UI_BINDING_CONFIG_VAR;
control.config_var = g_cvarEnabled;
add_control(panel, control);
control = UI_CONTROL_DESC_INIT;
control.kind = UI_CONTROL_BUTTON;
control.label = "Open Controls";
control.on_pressed = on_open_controls;
add_control(panel, control);
return MOD_OK;
}
ModResult register_bool_option(
const char* name, bool defaultValue, ConfigVarHandle& outHandle, ModError* error) {
ConfigVarDesc cvarDesc = CONFIG_VAR_DESC_INIT;
cvarDesc.name = name;
cvarDesc.type = CONFIG_VAR_BOOL;
cvarDesc.default_bool = defaultValue;
if (svc_config->register_var(mod_ctx, &cvarDesc, &outHandle) != MOD_OK) {
return dusk::mods::set_error(error, MOD_ERROR, "failed to register AO option");
}
return MOD_OK;
}
ModResult register_int_option(
const char* name, int64_t defaultValue, ConfigVarHandle& outHandle, ModError* error) {
ConfigVarDesc cvarDesc = CONFIG_VAR_DESC_INIT;
cvarDesc.name = name;
cvarDesc.type = CONFIG_VAR_INT;
cvarDesc.default_int = defaultValue;
if (svc_config->register_var(mod_ctx, &cvarDesc, &outHandle) != MOD_OK) {
return dusk::mods::set_error(error, MOD_ERROR, "failed to register AO option");
}
return MOD_OK;
}
} // namespace
extern "C" {
MOD_EXPORT ModResult mod_initialize(ModError* error) {
ModResult result = svc_resource->load(mod_ctx, "preprocess_depth.wgsl", &g_preprocessSource);
if (result == MOD_OK) {
result = svc_resource->load(mod_ctx, "gtao.wgsl", &g_gtaoSource);
}
if (result == MOD_OK) {
result = svc_resource->load(mod_ctx, "denoise.wgsl", &g_denoiseSource);
}
if (result == MOD_OK) {
result = svc_resource->load(mod_ctx, "composite.wgsl", &g_compositeSource);
}
if (result != MOD_OK) {
return dusk::mods::set_error(error, result, "failed to load AO shaders");
}
result = register_bool_option("effectEnabled", false, g_cvarEnabled, error);
if (result != MOD_OK) {
return result;
}
result = register_int_option("quality", 2, g_cvarQuality, error);
if (result != MOD_OK) {
return result;
}
result = register_int_option("radius", 70, g_cvarRadius, error);
if (result != MOD_OK) {
return result;
}
result = register_int_option("intensity", 100, g_cvarIntensity, error);
if (result != MOD_OK) {
return result;
}
result = register_bool_option("halfRes", false, g_cvarHalfRes, error);
if (result != MOD_OK) {
return result;
}
result = register_int_option("debugMode", 0, g_cvarDebugView, error);
if (result != MOD_OK) {
return result;
}
if (svc_gfx->get_device_info(mod_ctx, &g_deviceInfo) != MOD_OK) {
return dusk::mods::set_error(error, MOD_ERROR, "failed to query device info");
}
if (!build_compute_pipeline("AO preprocess depth", g_preprocessSource, "preprocess_depth",
g_preprocessPipeline, g_preprocessLayout) ||
!build_compute_pipeline("AO downsample mip4", g_preprocessSource, "downsample_mip4",
g_mip4Pipeline, g_mip4Layout) ||
!build_compute_pipeline("AO gtao", g_gtaoSource, "gtao", g_gtaoPipeline, g_gtaoLayout) ||
!build_compute_pipeline(
"AO denoise", g_denoiseSource, "spatial_denoise", g_denoisePipeline, g_denoiseLayout))
{
return dusk::mods::set_error(error, MOD_ERROR, "failed to create AO compute pipelines");
}
if (!build_composite_pipeline(true, g_compositePipeline, g_compositeLayout) ||
!build_composite_pipeline(false, g_compositeDebugPipeline, g_compositeDebugLayout))
{
return dusk::mods::set_error(error, MOD_ERROR, "failed to create AO composite pipeline");
}
if (!build_hilbert_lut()) {
return dusk::mods::set_error(error, MOD_ERROR, "failed to create AO noise LUT");
}
GfxComputeTypeDesc computeDesc = GFX_COMPUTE_TYPE_DESC_INIT;
computeDesc.label = "AO chain";
computeDesc.callback = on_compute;
if (svc_gfx->register_compute_type(mod_ctx, &computeDesc, &g_computeType) != MOD_OK) {
return dusk::mods::set_error(error, MOD_ERROR, "failed to register compute type");
}
GfxDrawTypeDesc drawDesc = GFX_DRAW_TYPE_DESC_INIT;
drawDesc.label = "AO composite";
drawDesc.draw = on_draw;
if (svc_gfx->register_draw_type(mod_ctx, &drawDesc, &g_drawType) != MOD_OK) {
return dusk::mods::set_error(error, MOD_ERROR, "failed to register draw type");
}
GfxStageHookDesc stageDesc = GFX_STAGE_HOOK_DESC_INIT;
stageDesc.callback = on_scene_after_opaque;
if (svc_gfx->register_stage_hook(
mod_ctx, GFX_STAGE_SCENE_AFTER_OPAQUE, &stageDesc, &g_afterOpaqueHook) != MOD_OK)
{
return dusk::mods::set_error(error, MOD_ERROR, "failed to register stage hook");
}
UiModsPanelDesc panelDesc = UI_MODS_PANEL_DESC_INIT;
panelDesc.build = build_panel;
svc_ui->register_mods_panel(mod_ctx, &panelDesc);
svc_log->info(mod_ctx, "ao_mod ready");
return MOD_OK;
}
MOD_EXPORT ModResult mod_update(ModError*) {
if (!g_loggedChain && g_chainExecuted.load(std::memory_order_acquire)) {
g_loggedChain = true;
svc_log->info(mod_ctx, "AO chain executed OK");
}
return MOD_OK;
}
MOD_EXPORT ModResult mod_shutdown(ModError*) {
svc_resource->free(mod_ctx, &g_preprocessSource);
svc_resource->free(mod_ctx, &g_gtaoSource);
svc_resource->free(mod_ctx, &g_denoiseSource);
svc_resource->free(mod_ctx, &g_compositeSource);
release_targets(g_targets);
for (auto& retired : g_retiredTargets) {
release_targets(retired.targets);
}
g_retiredTargets.clear();
const auto releasePipeline = [](WGPUComputePipeline& pipeline) {
if (pipeline != nullptr) {
wgpuComputePipelineRelease(pipeline);
pipeline = nullptr;
}
};
const auto releaseLayout = [](WGPUBindGroupLayout& layout) {
if (layout != nullptr) {
wgpuBindGroupLayoutRelease(layout);
layout = nullptr;
}
};
releasePipeline(g_preprocessPipeline);
releasePipeline(g_mip4Pipeline);
releasePipeline(g_gtaoPipeline);
releasePipeline(g_denoisePipeline);
releaseLayout(g_preprocessLayout);
releaseLayout(g_mip4Layout);
releaseLayout(g_gtaoLayout);
releaseLayout(g_denoiseLayout);
if (g_compositePipeline != nullptr) {
wgpuRenderPipelineRelease(g_compositePipeline);
g_compositePipeline = nullptr;
}
if (g_compositeDebugPipeline != nullptr) {
wgpuRenderPipelineRelease(g_compositeDebugPipeline);
g_compositeDebugPipeline = nullptr;
}
releaseLayout(g_compositeLayout);
releaseLayout(g_compositeDebugLayout);
if (g_hilbertLutView != nullptr) {
wgpuTextureViewRelease(g_hilbertLutView);
g_hilbertLutView = nullptr;
}
if (g_hilbertLut != nullptr) {
wgpuTextureRelease(g_hilbertLut);
g_hilbertLut = nullptr;
}
g_cvarEnabled = g_cvarQuality = g_cvarRadius = g_cvarIntensity = 0;
g_cvarHalfRes = g_cvarDebugView = 0;
g_computeType = g_drawType = 0;
g_afterOpaqueHook = 0;
g_controlsWindow = 0;
return MOD_OK;
}
}