ddgi_relocate.slang

#include "shared.h"
#include "renderer_types.h"

[[vk::push_constant]]
DdgiRelocatePushConstants pc;

static const f32 DDGI_MISS_DISTANCE = 1e27f;

[shader("compute")]
[numthreads(64, 1, 1)]
void csMain(uvec3 dtid: SV_DispatchThreadID) {
    uint probe_idx = dtid.x;
    FrameUBO frame = *pc.frame;
    DdgiGridParams grid = frame.ddgi_grid;
    DdgiConfig cfg = frame.ddgi_cfg;

    uvec3 dims = uvec3(grid.dims.xyz);
    if (probe_idx >= dims.x * dims.y * dims.z)
        return;

    // load current offset (denormalize to world) and state
    vec4 ps = pc.probe_state[probe_idx].offset_active;
    vec3 off = ps.xyz * grid.spacing.xyz; // denormalize to world-space
    f32 state = ps.w;

    // scan fixed rays for closest/farthest frontface and closest backface
    uint backface_hits = 0;
    int closest_backface_idx = -1;
    int closest_frontface_idx = -1;
    int farthest_frontface_idx = -1;
    f32 closest_backface_dist = 1e27f;
    f32 closest_frontface_dist = 1e27f;
    f32 farthest_frontface_dist = 0.0f;

    uint numRays = FIXED_RAY_COUNT;
    for (uint r = 0; r < numRays; r++) {
        vec4 rad_dist = pc.samples_buffer[probe_idx * RAYS_PER_PROBE + r].radiance_distance;

        if (abs(rad_dist.w) >= DDGI_MISS_DISTANCE)
            continue;

        if (rad_dist.w < 0.0f) {
            // backface hit
            backface_hits++;

            // backfaces store -dist * 0.2, restore with * -5
            f32 hitDist = rad_dist.w * -5.0f;
            if (hitDist < closest_backface_dist) {
                closest_backface_dist = hitDist;
                closest_backface_idx = (int)r;
            }
        } else {
            // frontface hit
            if (rad_dist.w < closest_frontface_dist) {
                closest_frontface_dist = rad_dist.w;
                closest_frontface_idx = (int)r;
            }
            if (rad_dist.w > farthest_frontface_dist) {
                farthest_frontface_dist = rad_dist.w;
                farthest_frontface_idx = (int)r;
            }
        }
    }

    vec3 fullOff = vec3(1e27f, 1e27f, 1e27f);

    // probe inside geometry, push out along closest backface dir
    if (closest_backface_idx != -1 && (f32(backface_hits) / f32(numRays)) > cfg.fixed_ray_backface_threshold) {
        vec3 closestBackfaceDir = pc.samples_buffer[probe_idx * RAYS_PER_PROBE + closest_backface_idx].direction.xyz;
        fullOff = off + normalize(closestBackfaceDir) * closest_backface_dist * (cfg.probe_distance_scale + 1.0f);
    }
    // probe too close to surface, nudge away
    else if (closest_frontface_dist < cfg.probe_min_frontface_distance) {
        vec3 closestDir = pc.samples_buffer[probe_idx * RAYS_PER_PROBE + closest_frontface_idx].direction.xyz;
        vec3 farthestDir = pc.samples_buffer[probe_idx * RAYS_PER_PROBE + farthest_frontface_idx].direction.xyz;

        // move if closest and farthest directions disagree, ensures we don't push through geometry
        if (dot(closestDir, farthestDir) <= 0.0f) {
            fullOff = off + normalize(farthestDir) * cfg.probe_distance_scale;
        }
    }
    // probe far from everything, drift back to zero offset
    else if (closest_frontface_dist > cfg.probe_min_frontface_distance + cfg.probe_distance_scale) {
        f32 moveBackMargin = min(closest_frontface_dist - cfg.probe_min_frontface_distance, length(off));
        vec3 moveBackDirection = normalize(-off);
        fullOff = off + moveBackMargin * moveBackDirection;
    }

    // keep probe within 45% of its grid cell
    vec3 normalizedOff = fullOff / grid.spacing.xyz;
    if (dot(normalizedOff, normalizedOff) < 0.2025f) {
        off = fullOff;
    }

    // store normalized offset, state (.w) written by classify
    pc.probe_state[probe_idx].offset_active = vec4(off / grid.spacing.xyz, state);
}