GPU-accelerated UI toolkit (Vulkan)
git clone https://git.lucas.co/cce-ui.git
src/vk/rt_common.wgsl (6.4K)
1 // rt_common.wgsl — the path tracer's shared core (RT-renderer phases 2+4).
2 //
3 // Everything except the trace call: params, scene/material buffers,
4 // accumulation, RNG, sky, sampling, and cs_main. Binding 1 and
5 // `intersect_scene` come from whichever tier file is concatenated after
6 // this one at pipeline creation:
7 // - rt_bvh.wgsl — tier 1: a CPU-built BVH traversed in compute; runs
8 // on any device, no VK_KHR_ray_* required.
9 // - rt_query.wgsl — tier 2: hardware ray queries against a driver-built
10 // TLAS (VK_KHR_ray_query), engaging RT cores.
11 // One dispatch adds `spp` samples per pixel into the accumulation buffer
12 // (progressive refinement); the running mean is tone-mapped (clamped
13 // linear) into `out_img`, which the stage blits into the backdrop pane.
14
15 struct Params {
16 // Inverse of the raster path's proj*view*model: unprojects wgpu-style NDC
17 // (y up, z in [0,1]) into mesh space, so rays live in the same space as
18 // the triangles fed to `set_rt_scene`.
19 inv_mvp: mat4x4<f32>,
20 width: u32,
21 height: u32,
22 sample_index: u32,
23 max_bounces: u32,
24 // Samples per dispatch: 1 for the interactive viewport (one refinement
25 // step per frame), higher for offscreen/thumbnail rendering so a whole
26 // image needs only a few submits.
27 spp: u32,
28 _pad0: u32,
29 _pad1: u32,
30 _pad2: u32,
31 }
32
33 @group(0) @binding(0) var<uniform> params: Params;
34
35 // Binding 1 belongs to the tier file: the BVH node buffer (tier 1) or the
36 // acceleration structure (tier 2).
37
38 // Positions in xyz; p0.w carries the material index (bitcast).
39 struct Tri {
40 p0: vec4<f32>,
41 p1: vec4<f32>,
42 p2: vec4<f32>,
43 }
44 @group(0) @binding(2) var<storage, read> tris: array<Tri>;
45
46 struct Material {
47 albedo: vec4<f32>,
48 emission: vec4<f32>,
49 }
50 @group(0) @binding(3) var<storage, read> materials: array<Material>;
51
52 // One vec4 per pixel: rgb = radiance sum, a = sample count.
53 @group(0) @binding(4) var<storage, read_write> accum: array<vec4<f32>>;
54
55 @group(0) @binding(5) var out_img: texture_storage_2d<rgba8unorm, write>;
56
57 // Primary-hit features for the denoiser (rt_denoise.wgsl), two vec4s per
58 // pixel: [2i] = (shading normal, hit t — 1e30 for sky), [2i+1] = (albedo, 0).
59 @group(0) @binding(6) var<storage, read_write> features: array<vec4<f32>>;
60
61 // PCG (O'Neill) — one u32 of state per path, advanced per draw.
62 fn rand(state: ptr<function, u32>) -> f32 {
63 var s = *state * 747796405u + 2891336453u;
64 *state = s;
65 let word = ((s >> ((s >> 28u) + 4u)) ^ s) * 277803737u;
66 return f32((word >> 22u) ^ word) * (1.0 / 4294967295.0);
67 }
68
69 // The tier boundary: whichever tier file follows provides
70 // fn intersect_scene(ro: vec3<f32>, rd: vec3<f32>) -> HitInfo
71 struct HitInfo {
72 t: f32,
73 tri: u32,
74 }
75
76 // A soft studio sky: vertical gradient plus one warm key light. This is the
77 // only light source until emissive geometry shows up in scenes.
78 fn sky(rd: vec3<f32>) -> vec3<f32> {
79 let t = clamp(rd.y * 0.5 + 0.5, 0.0, 1.0);
80 var s = mix(vec3<f32>(0.32, 0.31, 0.35), vec3<f32>(0.72, 0.82, 0.98), t);
81 let sun = normalize(vec3<f32>(0.45, 0.75, 0.35));
82 s = s + vec3<f32>(1.0, 0.95, 0.85) * pow(max(dot(rd, sun), 0.0), 48.0) * 8.0;
83 return s;
84 }
85
86 fn cosine_dir(n: vec3<f32>, r1: f32, r2: f32) -> vec3<f32> {
87 let a = 6.28318530718 * r1;
88 let r = sqrt(r2);
89 var up = vec3<f32>(1.0, 0.0, 0.0);
90 if abs(n.x) > 0.5 {
91 up = vec3<f32>(0.0, 1.0, 0.0);
92 }
93 let tangent = normalize(cross(n, up));
94 let bitangent = cross(n, tangent);
95 return normalize(
96 tangent * (r * cos(a)) + bitangent * (r * sin(a)) + n * sqrt(max(0.0, 1.0 - r2)),
97 );
98 }
99
100 @compute @workgroup_size(8, 8)
101 fn cs_main(@builtin(global_invocation_id) gid: vec3<u32>) {
102 if gid.x >= params.width || gid.y >= params.height {
103 return;
104 }
105 let idx = gid.y * params.width + gid.x;
106
107 var total = vec3<f32>(0.0);
108 for (var s: u32 = 0u; s < params.spp; s = s + 1u) {
109 var rng: u32 = (idx * 9781u) ^ ((params.sample_index + s) * 26699u) ^ 0x9e3779b9u;
110
111 // Jittered primary ray, unprojected through inv_mvp (NDC y up, z 0..1).
112 let jx = rand(&rng);
113 let jy = rand(&rng);
114 let ndc_x = (f32(gid.x) + jx) / f32(params.width) * 2.0 - 1.0;
115 let ndc_y = 1.0 - (f32(gid.y) + jy) / f32(params.height) * 2.0;
116 let p_near = params.inv_mvp * vec4<f32>(ndc_x, ndc_y, 0.0, 1.0);
117 let p_far = params.inv_mvp * vec4<f32>(ndc_x, ndc_y, 1.0, 1.0);
118 var ro = p_near.xyz / p_near.w;
119 var rd = normalize(p_far.xyz / p_far.w - ro);
120
121 var radiance = vec3<f32>(0.0);
122 var throughput = vec3<f32>(1.0);
123 for (var bounce: u32 = 0u; bounce < params.max_bounces; bounce = bounce + 1u) {
124 let hit = intersect_scene(ro, rd);
125 if hit.t >= 1e30 {
126 if s == 0u && bounce == 0u {
127 features[2u * idx] = vec4<f32>(0.0, 0.0, 0.0, 1e30);
128 features[2u * idx + 1u] = vec4<f32>(1.0, 1.0, 1.0, 0.0);
129 }
130 radiance = radiance + throughput * sky(rd);
131 break;
132 }
133 let tri = tris[hit.tri];
134 let mat = materials[bitcast<u32>(tri.p0.w)];
135 radiance = radiance + throughput * mat.emission.rgb;
136 var n = normalize(cross(tri.p1.xyz - tri.p0.xyz, tri.p2.xyz - tri.p0.xyz));
137 if dot(n, rd) > 0.0 {
138 n = -n;
139 }
140 if s == 0u && bounce == 0u {
141 features[2u * idx] = vec4<f32>(n, hit.t);
142 features[2u * idx + 1u] = vec4<f32>(mat.albedo.rgb, 0.0);
143 }
144 throughput = throughput * mat.albedo.rgb;
145 ro = ro + rd * hit.t + n * 1e-4;
146 rd = cosine_dir(n, rand(&rng), rand(&rng));
147 }
148 // Firefly clamp: rare sun-spike paths otherwise leave speckles the
149 // variance can't average out (and the denoiser's edge-stopping
150 // weights deliberately refuse to smear). Slight energy loss on
151 // extreme highlights, big variance win.
152 total = total + min(radiance, vec3<f32>(4.0));
153 }
154
155 var acc = accum[idx];
156 if params.sample_index == 0u {
157 acc = vec4<f32>(0.0);
158 }
159 acc = acc + vec4<f32>(total, f32(params.spp));
160 accum[idx] = acc;
161 let color = acc.rgb / max(acc.a, 1.0);
162 textureStore(
163 out_img,
164 vec2<i32>(i32(gid.x), i32(gid.y)),
165 vec4<f32>(clamp(color, vec3<f32>(0.0), vec3<f32>(1.0)), 1.0),
166 );
167 }