GPU-accelerated UI toolkit (Vulkan)
git clone https://git.lucas.co/cce-ui.git
src/vk/rt_denoise.wgsl (4.2K)
1 // rt_denoise.wgsl — edge-avoiding à-trous wavelet denoiser (Dammertz et al.)
2 // for the path tracer's progressive output.
3 //
4 // Runs as a few compute iterations after each trace dispatch: a 5×5 B3-spline
5 // kernel dilated by `step` (1, 2, 4, …), with edge-stopping weights from the
6 // tracer's primary-hit features (normal + depth) and color similarity. The
7 // color sigma scales with 1/sqrt(sample count), so the filter is strong on
8 // 1-spp camera-drag frames and fades toward identity as the accumulation
9 // converges — no pop when it effectively "turns off".
10 //
11 // Iteration wiring (single pipeline, per-iteration dynamic uniform):
12 // first != 0 — source is the accumulation mean (accum.rgb / accum.a);
13 // otherwise the `src` ping-pong buffer.
14 // last != 0 — the result is also tone-mapped into `out_img` (replacing
15 // the tracer's own write); it is always written to `dst`.
16
17 struct DenoiseParams {
18 width: u32,
19 height: u32,
20 step: u32,
21 first: u32,
22 last: u32,
23 // 1/sqrt(accumulated samples): scales the color edge-stopping sigma.
24 inv_sqrt_n: f32,
25 _pad0: u32,
26 _pad1: u32,
27 }
28
29 @group(0) @binding(0) var<uniform> dp: DenoiseParams;
30 @group(0) @binding(1) var<storage, read> accum: array<vec4<f32>>;
31 @group(0) @binding(2) var<storage, read> features: array<vec4<f32>>;
32 @group(0) @binding(3) var<storage, read> src: array<vec4<f32>>;
33 @group(0) @binding(4) var<storage, read_write> dst: array<vec4<f32>>;
34 @group(0) @binding(5) var out_img: texture_storage_2d<rgba8unorm, write>;
35
36 fn load_color(idx: u32) -> vec3<f32> {
37 if dp.first != 0u {
38 let a = accum[idx];
39 return a.rgb / max(a.a, 1.0);
40 }
41 return src[idx].rgb;
42 }
43
44 fn luminance(c: vec3<f32>) -> f32 {
45 return dot(c, vec3<f32>(0.2126, 0.7152, 0.0722));
46 }
47
48 // B3 spline: center, ±1, ±2.
49 fn kernel_w(k: i32) -> f32 {
50 if k == 0 {
51 return 0.375;
52 }
53 if k == 1 {
54 return 0.25;
55 }
56 return 0.0625;
57 }
58
59 @compute @workgroup_size(8, 8)
60 fn cs_denoise(@builtin(global_invocation_id) gid: vec3<u32>) {
61 if gid.x >= dp.width || gid.y >= dp.height {
62 return;
63 }
64 let idx = gid.y * dp.width + gid.x;
65 let center_c = load_color(idx);
66 let center_f = features[2u * idx];
67 let center_n = center_f.xyz;
68 let center_t = center_f.w;
69 let center_l = luminance(center_c);
70
71 // Sky pixels (t = 1e30) are already noise-free; pass them through rather
72 // than letting the huge depth deltas produce all-zero weights.
73 if center_t >= 1e30 {
74 dst[idx] = vec4<f32>(center_c, 1.0);
75 if dp.last != 0u {
76 textureStore(
77 out_img,
78 vec2<i32>(i32(gid.x), i32(gid.y)),
79 vec4<f32>(clamp(center_c, vec3<f32>(0.0), vec3<f32>(1.0)), 1.0),
80 );
81 }
82 return;
83 }
84
85 let sigma_c = max(0.85 * dp.inv_sqrt_n, 1e-3);
86 let sigma_n = 0.15;
87 let sigma_t = 0.10 * max(center_t, 0.1);
88
89 var sum = vec3<f32>(0.0);
90 var weight_sum = 0.0;
91 for (var dy: i32 = -2; dy <= 2; dy = dy + 1) {
92 for (var dx: i32 = -2; dx <= 2; dx = dx + 1) {
93 let x = i32(gid.x) + dx * i32(dp.step);
94 let y = i32(gid.y) + dy * i32(dp.step);
95 if x < 0 || y < 0 || x >= i32(dp.width) || y >= i32(dp.height) {
96 continue;
97 }
98 let j = u32(y) * dp.width + u32(x);
99 let f = features[2u * j];
100 if f.w >= 1e30 {
101 continue; // never mix sky into surfaces
102 }
103 let c = load_color(j);
104
105 let dn = center_n - f.xyz;
106 let w_n = exp(-dot(dn, dn) / (sigma_n * sigma_n));
107 let dt = (center_t - f.w) / sigma_t;
108 let w_t = exp(-dt * dt);
109 let dl = (center_l - luminance(c)) / sigma_c;
110 let w_c = exp(-dl * dl);
111 let w = kernel_w(abs(dx)) * kernel_w(abs(dy)) * w_n * w_t * w_c;
112
113 sum = sum + c * w;
114 weight_sum = weight_sum + w;
115 }
116 }
117 let result = sum / max(weight_sum, 1e-6);
118 dst[idx] = vec4<f32>(result, 1.0);
119 if dp.last != 0u {
120 textureStore(
121 out_img,
122 vec2<i32>(i32(gid.x), i32(gid.y)),
123 vec4<f32>(clamp(result, vec3<f32>(0.0), vec3<f32>(1.0)), 1.0),
124 );
125 }
126 }