Skip to main content

RESOLVE_WGSL

Constant RESOLVE_WGSL 

Source
pub const RESOLVE_WGSL: &str = r#"
struct Uniforms {
    width: u32,
    height: u32,
    flags: u32,
    exposure: f32,
    history_exposure: f32,
    depth_threshold: f32,
    normal_min_dot: f32,
    variance_gamma: f32,
    history_weight: f32,
    history_width: u32,
    history_height: u32,
}

@group(0) @binding(0) var<uniform> u: Uniforms;
@group(0) @binding(1) var colour_tex: texture_2d<f32>;
@group(0) @binding(2) var depth_tex: texture_2d<f32>;
@group(0) @binding(3) var motion_tex: texture_2d<f32>;
@group(0) @binding(4) var id_tex: texture_2d<u32>;
@group(0) @binding(5) var normal_tex: texture_2d<f32>;
@group(0) @binding(6) var reactive_tex: texture_2d<f32>;
@group(0) @binding(7) var history_colour: texture_2d<f32>;
@group(0) @binding(8) var history_depth: texture_2d<f32>;
@group(0) @binding(9) var history_id: texture_2d<u32>;
@group(0) @binding(10) var history_normal: texture_2d<f32>;
@group(0) @binding(11) var out_colour: texture_storage_2d<rgba16float, write>;
@group(0) @binding(12) var out_depth: texture_storage_2d<r32float, write>;
@group(0) @binding(13) var out_id: texture_storage_2d<r32uint, write>;
@group(0) @binding(14) var out_normal: texture_storage_2d<rgba8snorm, write>;

fn rgb_to_ycocg(c: vec3<f32>) -> vec3<f32> {
    return vec3<f32>(
        0.25 * c.r + 0.5 * c.g + 0.25 * c.b,
        0.5 * c.r - 0.5 * c.b,
        -0.25 * c.r + 0.5 * c.g - 0.25 * c.b,
    );
}

fn ycocg_to_rgb(c: vec3<f32>) -> vec3<f32> {
    return vec3<f32>(c.x + c.y - c.z, c.x + c.z, c.x - c.y - c.z);
}

fn outside(uv: vec2<f32>) -> bool {
    return uv.x < 0.0 || uv.y < 0.0 || uv.x >= 1.0 || uv.y >= 1.0;
}

fn clamp_pixel(p: vec2<i32>, size: vec2<i32>) -> vec2<i32> {
    return clamp(p, vec2<i32>(0), size - vec2<i32>(1));
}

fn load3(tex: texture_2d<f32>, p: vec2<i32>, size: vec2<i32>) -> vec3<f32> {
    return textureLoad(tex, clamp_pixel(p, size), 0).rgb;
}

fn bilinear3(tex: texture_2d<f32>, uv: vec2<f32>, size: vec2<u32>) -> vec3<f32> {
    let s = vec2<i32>(size);
    let pf = uv * vec2<f32>(size) - vec2<f32>(0.5);
    let base = vec2<i32>(floor(pf));
    let f = pf - floor(pf);
    let c00 = load3(tex, base, s);
    let c10 = load3(tex, base + vec2<i32>(1, 0), s);
    let c01 = load3(tex, base + vec2<i32>(0, 1), s);
    let c11 = load3(tex, base + vec2<i32>(1, 1), s);
    return mix(mix(c00, c10, f.x), mix(c01, c11, f.x), f.y);
}

fn nearest(uv: vec2<f32>, size: vec2<u32>) -> vec2<i32> {
    return clamp_pixel(vec2<i32>(floor(uv * vec2<f32>(size))), vec2<i32>(size));
}

fn clip_aabb(lo: vec3<f32>, hi: vec3<f32>, p: vec3<f32>) -> vec3<f32> {
    let center = 0.5 * (hi + lo);
    let extent = max(0.5 * (hi - lo), vec3<f32>(1.0e-8));
    let delta = p - center;
    let unit = delta / extent;
    let m = max(abs(unit.x), max(abs(unit.y), abs(unit.z)));
    if (m > 1.0) {
        return center + delta / m;
    }
    return p;
}

@compute @workgroup_size(8, 8)
fn resolve(@builtin(global_invocation_id) gid: vec3<u32>) {
    let size = vec2<u32>(u.width, u.height);
    if (gid.x >= size.x || gid.y >= size.y) {
        return;
    }
    let pix = vec2<i32>(gid.xy);
    let uv = (vec2<f32>(gid.xy) + vec2<f32>(0.5)) / vec2<f32>(size);
    let current = textureLoad(colour_tex, pix, 0).rgb;
    let depth = textureLoad(depth_tex, pix, 0).r;
    let motion = textureLoad(motion_tex, pix, 0).rg;
    let reactive = clamp(textureLoad(reactive_tex, pix, 0).r, 0.0, 1.0);
    var pixel_id = 0u;
    if ((u.flags & 2u) != 0u) {
        pixel_id = textureLoad(id_tex, pix, 0).r;
    }
    var nrm = vec3<f32>(0.0);
    if ((u.flags & 4u) != 0u) {
        nrm = textureLoad(normal_tex, pix, 0).rgb;
    }

    var acc = vec3<f32>(0.0);
    var acc2 = vec3<f32>(0.0);
    let s = vec2<i32>(size);
    for (var oy = -1; oy <= 1; oy = oy + 1) {
        for (var ox = -1; ox <= 1; ox = ox + 1) {
            let texel = load3(colour_tex, pix + vec2<i32>(ox, oy), s) * u.exposure;
            let ycocg = rgb_to_ycocg(texel);
            acc = acc + ycocg;
            acc2 = acc2 + ycocg * ycocg;
        }
    }
    let mean = acc / 9.0;
    let sigma = sqrt(max(acc2 / 9.0 - mean * mean, vec3<f32>(0.0))) * u.variance_gamma;
    let lo = mean - sigma;
    let hi = mean + sigma;
    let current_e = current * u.exposure;
    var history_e = current_e;
    var weight = 0.0;
    if ((u.flags & 1u) != 0u) {
        let prev = uv + motion;
        if (!outside(prev)) {
            let stored = max(u.history_exposure, 1.0e-8);
            let history_size = vec2<u32>(u.history_width, u.history_height);
            let history_hdr = bilinear3(history_colour, prev, history_size);
            history_e = (history_hdr * stored) * (u.exposure / stored);
            var accept = true;
            let at = nearest(prev, history_size);
            let previous_depth = textureLoad(history_depth, at, 0).r;
            let scale = max(max(abs(depth), abs(previous_depth)), 1.0e-5);
            if (abs(depth - previous_depth) / scale > u.depth_threshold) {
                accept = false;
            }
            if ((u.flags & 2u) != 0u) {
                let previous_id = textureLoad(history_id, at, 0).r;
                if (previous_id != pixel_id) {
                    accept = false;
                }
            }
            if ((u.flags & 4u) != 0u) {
                let previous_normal = textureLoad(history_normal, at, 0).rgb;
                let current_len2 = dot(nrm, nrm);
                let previous_len2 = dot(previous_normal, previous_normal);
                if (current_len2 > 1.0e-12 && previous_len2 > 1.0e-12) {
                    let aligned = dot(nrm, previous_normal) * inverseSqrt(current_len2 * previous_len2);
                    if (aligned < u.normal_min_dot) {
                        accept = false;
                    }
                }
            }
            if (accept) {
                history_e = ycocg_to_rgb(clip_aabb(lo, hi, rgb_to_ycocg(history_e)));
                weight = u.history_weight * (1.0 - reactive);
            }
        }
    }
    let resolved_e = history_e * weight + current_e * (1.0 - weight);
    let resolved = resolved_e / u.exposure;
    textureStore(out_colour, pix, vec4<f32>(resolved, u.exposure));
    textureStore(out_depth, pix, vec4<f32>(depth, 0.0, 0.0, 0.0));
    textureStore(out_id, pix, vec4<u32>(pixel_id, 0u, 0u, 0u));
    textureStore(out_normal, pix, vec4<f32>(nrm, 0.0));
}
"#;