GPU-accelerated UI toolkit (Vulkan)
git clone https://git.lucas.co/cce-ui.git
src/draw/rt.rs (28.3K)
1 //! The path tracer's scene and the work done on the CPU for it, with no GPU
2 //! in it: the schema an app fills ([`RtTriangle`], [`RtMaterial`],
3 //! [`RtImage`], [`RtCamera`], [`RtEnvironment`]), the binned-SAH BVH built
4 //! over it, the buffers and parameter blocks laid out as `rt_common.wgsl` /
5 //! `rt_bvh.wgsl` / `rt_denoise.wgsl` read them, and the tracer's constants.
6 //! The Vulkan stage (`vk::rt`, with a hardware ray-query tier besides) and
7 //! the WebGPU one (`web::rt`, the compute tier) both trace from these.
8 //!
9 //! The tracer is plain compute: a CPU-built BVH traversed per pixel,
10 //! progressive accumulation of one sample a frame, an à-trous denoise over
11 //! the running mean, blitted into the backdrop's pane region in place of the
12 //! raster scene. Scene schema is internal by design — importers (OBJ/glTF)
13 //! belong in a loader that converts *into* [`RtTriangle`]/[`RtMaterial`].
14
15 /// One triangle of an RT scene, in the same space as the camera's `inv_mvp`
16 /// (for the designer: mesh space, the space `Vertex3D` positions live in).
17 #[derive(Debug, Clone, Copy, PartialEq)]
18 pub struct RtTriangle {
19 pub p0: [f32; 3],
20 pub p1: [f32; 3],
21 pub p2: [f32; 3],
22 /// Index into the material slice passed alongside.
23 pub material: u32,
24 }
25
26 /// Lambertian surface + optional emission, linear color (matching the raster
27 /// path, whose vertex colors land in the sRGB attachment as linear values).
28 #[derive(Debug, Clone, Copy, PartialEq)]
29 pub struct RtMaterial {
30 pub albedo: [f32; 3],
31 pub emission: [f32; 3],
32 }
33
34 /// An image standing in the traced scene: a quad that shows the image's
35 /// colour as it is — unlit, as the raster pass's `SceneImage` draws it — and
36 /// lets a ray through where the image is clear. A path ends on it, so to
37 /// the rest of the scene it is a light of its own colour. The tracer adds
38 /// the quad to the scene itself.
39 ///
40 /// The image is one the 2D pass already holds (an id from `upload_rgba`),
41 /// so a picture shown in the raster viewport costs the tracer nothing more.
42 #[derive(Debug, Clone, Copy, PartialEq)]
43 pub struct RtImage {
44 pub image: u32,
45 /// The quad's corners in the scene's space, in the image's own order:
46 /// top-left, top-right, bottom-right, bottom-left.
47 pub corners: [[f32; 3]; 4],
48 /// Alpha multiplier over the image's own.
49 pub opacity: f32,
50 }
51
52 /// [`RtImage`] for the headless tracer, which has no 2D pass to share an
53 /// image with and is handed the pixels: tightly packed sRGB RGBA8.
54 #[derive(Debug, Clone, Copy)]
55 pub struct RtImagePixels<'a> {
56 pub pixels: &'a [u8],
57 pub width: u32,
58 pub height: u32,
59 pub corners: [[f32; 3]; 4],
60 pub opacity: f32,
61 }
62
63 /// The full camera: the inverse of the raster path's `proj * view * model`.
64 /// Rays are unprojected from NDC through it, so any matrix stack that renders
65 /// the raster viewport drives the tracer unchanged.
66 #[derive(Debug, Clone, Copy, PartialEq)]
67 pub struct RtCamera {
68 pub inv_mvp: [[f32; 4]; 4],
69 }
70
71 // --- GPU layouts (must match rt.wgsl) ---
72
73 #[repr(C)]
74 #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
75 pub(crate) struct GpuTriangle {
76 pub(crate) p0: [f32; 4], // w = material index (bitcast)
77 pub(crate) p1: [f32; 4],
78 pub(crate) p2: [f32; 4],
79 }
80
81 #[repr(C)]
82 #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
83 pub(crate) struct GpuMaterial {
84 pub(crate) albedo: [f32; 4],
85 pub(crate) emission: [f32; 4],
86 }
87
88 #[repr(C)]
89 #[derive(Debug, Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
90 pub(crate) struct GpuBvhNode {
91 pub(crate) min: [f32; 3],
92 /// Leaf (`count > 0`): first triangle. Internal: left child; right = +1.
93 pub(crate) left_first: u32,
94 pub(crate) max: [f32; 3],
95 pub(crate) count: u32,
96 }
97
98 #[repr(C)]
99 #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
100 pub(crate) struct RtParams {
101 pub(crate) inv_mvp: [[f32; 4]; 4],
102 pub(crate) width: u32,
103 pub(crate) height: u32,
104 pub(crate) sample_index: u32,
105 pub(crate) max_bounces: u32,
106 pub(crate) spp: u32,
107 pub(crate) _pad: [u32; 3],
108 pub(crate) img_origin: [f32; 4],
109 pub(crate) img_u: [f32; 4],
110 pub(crate) img_v: [f32; 4],
111 pub(crate) background: [f32; 4],
112 // The environment, each xyz in a vec4 (w unused): toward the sun
113 // (unit), the sun's radiance, the sky overhead, the sky below.
114 pub(crate) sun_dir: [f32; 4],
115 pub(crate) sun_color: [f32; 4],
116 pub(crate) sky_zenith: [f32; 4],
117 pub(crate) sky_nadir: [f32; 4],
118 }
119
120 /// The traced scene's light: a sky that grades from `sky_nadir` straight
121 /// down to `sky_zenith` straight up, and a sun — a bright lobe toward
122 /// `sun_direction` of `sun_color` radiance. It is the tracer's ONLY light
123 /// (nothing in a scene emits unless its material does). Colours are linear
124 /// RGB and may exceed 1. [`Default`] is the studio sky the tracer always
125 /// had.
126 #[derive(Clone, Copy, Debug, PartialEq)]
127 pub struct RtEnvironment {
128 /// Toward the sun, world space; any length (zero is straight up).
129 pub sun_direction: [f32; 3],
130 pub sun_color: [f32; 3],
131 pub sky_zenith: [f32; 3],
132 pub sky_nadir: [f32; 3],
133 }
134
135 impl Default for RtEnvironment {
136 fn default() -> Self {
137 Self {
138 sun_direction: [0.45, 0.75, 0.35],
139 sun_color: [8.0, 7.6, 6.8],
140 sky_zenith: [0.72, 0.82, 0.98],
141 sky_nadir: [0.32, 0.31, 0.35],
142 }
143 }
144 }
145
146 impl RtEnvironment {
147 pub(crate) fn sun_unit(&self) -> [f32; 3] {
148 let v = glam::Vec3::from_array(self.sun_direction);
149 if v.length_squared() > 1e-12 { v.normalize().to_array() } else { [0.0, 1.0, 0.0] }
150 }
151 }
152
153 // --- BVH construction (binned SAH) ---
154
155 const BVH_BINS: usize = 8;
156 const BVH_LEAF_MAX: u32 = 4;
157
158 #[derive(Clone, Copy)]
159 struct Aabb {
160 min: [f32; 3],
161 max: [f32; 3],
162 }
163
164 impl Aabb {
165 const EMPTY: Aabb = Aabb { min: [f32::INFINITY; 3], max: [f32::NEG_INFINITY; 3] };
166
167 fn grow(&mut self, p: [f32; 3]) {
168 for a in 0..3 {
169 self.min[a] = self.min[a].min(p[a]);
170 self.max[a] = self.max[a].max(p[a]);
171 }
172 }
173
174 fn grow_aabb(&mut self, other: &Aabb) {
175 self.grow(other.min);
176 self.grow(other.max);
177 }
178
179 fn half_area(&self) -> f32 {
180 let dx = (self.max[0] - self.min[0]).max(0.0);
181 let dy = (self.max[1] - self.min[1]).max(0.0);
182 let dz = (self.max[2] - self.min[2]).max(0.0);
183 dx * dy + dy * dz + dz * dx
184 }
185 }
186
187 fn tri_aabb(t: &RtTriangle) -> Aabb {
188 let mut b = Aabb::EMPTY;
189 b.grow(t.p0);
190 b.grow(t.p1);
191 b.grow(t.p2);
192 b
193 }
194
195 fn tri_centroid(t: &RtTriangle) -> [f32; 3] {
196 let mut c = [0.0f32; 3];
197 for a in 0..3 {
198 c[a] = (t.p0[a] + t.p1[a] + t.p2[a]) / 3.0;
199 }
200 c
201 }
202
203 /// Build a BVH over `triangles`, reordering them so leaves reference
204 /// contiguous ranges. Returns the flat node array (empty input → empty vec).
205 pub(crate) fn build_bvh(triangles: &mut Vec<RtTriangle>) -> Vec<GpuBvhNode> {
206 if triangles.is_empty() {
207 return Vec::new();
208 }
209 let bounds: Vec<Aabb> = triangles.iter().map(tri_aabb).collect();
210 let centroids: Vec<[f32; 3]> = triangles.iter().map(tri_centroid).collect();
211 let mut order: Vec<u32> = (0..triangles.len() as u32).collect();
212
213 fn range_bounds(order: &[u32], bounds: &[Aabb]) -> Aabb {
214 let mut b = Aabb::EMPTY;
215 for &i in order {
216 b.grow_aabb(&bounds[i as usize]);
217 }
218 b
219 }
220
221 let mut nodes: Vec<GpuBvhNode> = Vec::with_capacity(triangles.len() * 2);
222 let root_bounds = range_bounds(&order, &bounds);
223 nodes.push(GpuBvhNode {
224 min: root_bounds.min,
225 left_first: 0,
226 max: root_bounds.max,
227 count: triangles.len() as u32,
228 });
229
230 // (node index, start, count) work list over `order`.
231 let mut work = vec![(0usize, 0usize, triangles.len())];
232 while let Some((node_idx, start, count)) = work.pop() {
233 if (count as u32) <= BVH_LEAF_MAX {
234 continue; // stays a leaf
235 }
236 let slice = &mut order[start..start + count];
237
238 // Centroid bounds pick the split axis.
239 let mut cb = Aabb::EMPTY;
240 for &i in slice.iter() {
241 cb.grow(centroids[i as usize]);
242 }
243 let mut axis = 0;
244 let mut extent = 0.0f32;
245 for a in 0..3 {
246 let e = cb.max[a] - cb.min[a];
247 if e > extent {
248 extent = e;
249 axis = a;
250 }
251 }
252
253 let mut split_at = None;
254 if extent > 1e-12 {
255 // Binned SAH along `axis`.
256 let scale = BVH_BINS as f32 / extent;
257 let bin_of = |i: u32| -> usize {
258 (((centroids[i as usize][axis] - cb.min[axis]) * scale) as usize)
259 .min(BVH_BINS - 1)
260 };
261 let mut bin_bounds = [Aabb::EMPTY; BVH_BINS];
262 let mut bin_counts = [0usize; BVH_BINS];
263 for &i in slice.iter() {
264 let b = bin_of(i);
265 bin_counts[b] += 1;
266 bin_bounds[b].grow_aabb(&bounds[i as usize]);
267 }
268 // Cost of each of the BINS-1 split planes.
269 let mut best_cost = f32::INFINITY;
270 let mut best_plane = 0usize;
271 for plane in 1..BVH_BINS {
272 let (mut lb, mut rb) = (Aabb::EMPTY, Aabb::EMPTY);
273 let (mut lc, mut rc) = (0usize, 0usize);
274 for b in 0..plane {
275 lb.grow_aabb(&bin_bounds[b]);
276 lc += bin_counts[b];
277 }
278 for b in plane..BVH_BINS {
279 rb.grow_aabb(&bin_bounds[b]);
280 rc += bin_counts[b];
281 }
282 if lc == 0 || rc == 0 {
283 continue;
284 }
285 let cost = lb.half_area() * lc as f32 + rb.half_area() * rc as f32;
286 if cost < best_cost {
287 best_cost = cost;
288 best_plane = plane;
289 }
290 }
291 if best_plane > 0 {
292 let mut mid = 0usize;
293 for k in 0..count {
294 if bin_of(slice[k]) < best_plane {
295 slice.swap(k, mid);
296 mid += 1;
297 }
298 }
299 if mid > 0 && mid < count {
300 split_at = Some(mid);
301 }
302 }
303 }
304 // Degenerate centroids or a one-sided SAH result: median split keeps
305 // the tree balanced instead of forcing a giant leaf.
306 let mid = split_at.unwrap_or(count / 2);
307
308 let left_bounds = range_bounds(&slice[..mid], &bounds);
309 let right_bounds = range_bounds(&slice[mid..], &bounds);
310 let left_idx = nodes.len();
311 nodes.push(GpuBvhNode {
312 min: left_bounds.min,
313 left_first: (start) as u32,
314 max: left_bounds.max,
315 count: mid as u32,
316 });
317 nodes.push(GpuBvhNode {
318 min: right_bounds.min,
319 left_first: (start + mid) as u32,
320 max: right_bounds.max,
321 count: (count - mid) as u32,
322 });
323 nodes[node_idx].left_first = left_idx as u32;
324 nodes[node_idx].count = 0;
325 work.push((left_idx, start, mid));
326 work.push((left_idx + 1, start + mid, count - mid));
327 }
328
329 // Apply the final order to the triangle array so leaf ranges are direct.
330 let reordered: Vec<RtTriangle> =
331 order.iter().map(|&i| triangles[i as usize]).collect();
332 *triangles = reordered;
333 nodes
334 }
335
336 // --- The Vulkan stage ---
337
338 pub(crate) const MAX_SAMPLES: u32 = 1024;
339 pub(crate) const MAX_BOUNCES: u32 = 4;
340 pub(crate) const WORKGROUP: u32 = 8;
341
342 #[repr(C)]
343 #[derive(Clone, Copy, bytemuck::Pod, bytemuck::Zeroable)]
344 pub(crate) struct DenoiseParams {
345 pub(crate) width: u32,
346 pub(crate) height: u32,
347 pub(crate) step: u32,
348 pub(crate) first: u32,
349 pub(crate) last: u32,
350 pub(crate) inv_sqrt_n: f32,
351 pub(crate) _pad: [u32; 2],
352 }
353
354 pub(crate) const DENOISE_ITERATIONS: usize = 3; // à-trous steps 1, 2, 4
355
356 /// A scene as the tracer's buffers hold it.
357 #[derive(Clone)]
358 pub(crate) struct PackedScene {
359 pub tris: Vec<GpuTriangle>,
360 pub materials: Vec<GpuMaterial>,
361 /// Empty unless `with_bvh` was asked for (the ray-query tier builds
362 /// its own structure).
363 pub nodes: Vec<GpuBvhNode>,
364 /// The BVH was built (`nodes` may still be empty: an empty scene).
365 pub has_bvh: bool,
366 }
367
368 impl PackedScene {
369 /// This scene with its BVH: itself when it has one, else one built
370 /// now over a copy of the triangles — the compute tier handed a scene
371 /// prepared for the ray-query one.
372 pub(crate) fn with_bvh(&self) -> std::borrow::Cow<'_, PackedScene> {
373 if self.has_bvh {
374 return std::borrow::Cow::Borrowed(self);
375 }
376 let xyz = |p: [f32; 4]| [p[0], p[1], p[2]];
377 let mut tris: Vec<RtTriangle> = self
378 .tris
379 .iter()
380 .map(|t| RtTriangle { p0: xyz(t.p0), p1: xyz(t.p1), p2: xyz(t.p2), material: t.p0[3].to_bits() })
381 .collect();
382 let nodes = build_bvh(&mut tris);
383 std::borrow::Cow::Owned(PackedScene {
384 tris: tris.iter().map(gpu_triangle).collect(),
385 materials: self.materials.clone(),
386 nodes,
387 has_bvh: true,
388 })
389 }
390 }
391
392 fn gpu_triangle(t: &RtTriangle) -> GpuTriangle {
393 GpuTriangle {
394 p0: [t.p0[0], t.p0[1], t.p0[2], f32::from_bits(t.material)],
395 p1: [t.p1[0], t.p1[1], t.p1[2], 0.0],
396 p2: [t.p2[0], t.p2[1], t.p2[2], 0.0],
397 }
398 }
399
400 /// Lay a scene out for the tracer. An image joins it as a quad of two
401 /// triangles (corners in its own order, top-left first) under a material
402 /// of its own, marked textured (`albedo.w`), past the scene's materials —
403 /// or past the one stood in for a scene that names none — so every tier
404 /// meets it as it meets any triangle, and a scene that is an image alone is
405 /// not an empty one. `with_bvh` builds the BVH, reordering the triangles.
406 pub(crate) fn pack_scene(
407 mut tris: Vec<RtTriangle>,
408 materials: &[RtMaterial],
409 image_corners: Option<[[f32; 3]; 4]>,
410 with_bvh: bool,
411 ) -> PackedScene {
412 let image_material = materials.len().max(1) as u32;
413 if let Some([tl, tr, br, bl]) = image_corners {
414 tris.push(RtTriangle { p0: tl, p1: bl, p2: tr, material: image_material });
415 tris.push(RtTriangle { p0: tr, p1: bl, p2: br, material: image_material });
416 }
417 let nodes = if with_bvh { build_bvh(&mut tris) } else { Vec::new() };
418 let gpu_tris = tris.iter().map(gpu_triangle).collect();
419 drop(tris);
420 let mut gpu_mats: Vec<GpuMaterial> = if materials.is_empty() {
421 vec![GpuMaterial { albedo: [0.8, 0.8, 0.8, 0.0], emission: [0.0; 4] }]
422 } else {
423 materials
424 .iter()
425 .map(|m| GpuMaterial {
426 albedo: [m.albedo[0], m.albedo[1], m.albedo[2], 0.0],
427 emission: [m.emission[0], m.emission[1], m.emission[2], 0.0],
428 })
429 .collect()
430 };
431 if image_corners.is_some() {
432 // albedo.w marks it textured: the shader takes the colour from the image.
433 gpu_mats.push(GpuMaterial { albedo: [1.0, 1.0, 1.0, 1.0], emission: [0.0; 4] });
434 }
435 PackedScene { tris: gpu_tris, materials: gpu_mats, nodes, has_bvh: with_bvh }
436 }
437
438 /// A traced scene with the CPU's share of the work already done: the
439 /// triangles packed into the tracer's buffer layouts and, for the compute
440 /// tier, the BVH built over them. Building one is plain CPU work with no
441 /// renderer in it — seconds for millions of triangles — so an app builds it
442 /// on a worker thread and hands it to
443 /// [`Stage3D::set_rt_scene_prepared`](crate::draw::scene::Stage3D::set_rt_scene_prepared)
444 /// on the UI thread, which only uploads it. `set_rt_scene` does both in
445 /// one call, on whatever thread calls it.
446 ///
447 /// It is not consumed by the upload: keep it to upload again into a
448 /// renderer rebuilt after a reconnect.
449 #[derive(Clone)]
450 pub struct PreparedRtScene {
451 pub(crate) packed: PackedScene,
452 pub(crate) image: Option<RtImage>,
453 }
454
455 impl PreparedRtScene {
456 /// Prepare `triangles` (taken, so millions of them are not copied) and
457 /// `materials`, with `image` standing in the scene as
458 /// `set_rt_scene_with_image` takes it. `with_bvh` builds the BVH: pass
459 /// the renderer's [`Stage3D::rt_needs_bvh`](crate::draw::scene::Stage3D::rt_needs_bvh),
460 /// asked on the UI thread before the work is sent off. A scene prepared
461 /// without one still traces on a renderer that needs it — the upload
462 /// builds it then, on the UI thread, as `set_rt_scene` would.
463 pub fn new(
464 triangles: Vec<RtTriangle>,
465 materials: &[RtMaterial],
466 image: Option<RtImage>,
467 with_bvh: bool,
468 ) -> Self {
469 PreparedRtScene { packed: pack_scene(triangles, materials, image.map(|i| i.corners), with_bvh), image }
470 }
471
472 /// Triangles the tracer holds, an image's quad included.
473 pub fn triangle_count(&self) -> usize {
474 self.packed.tris.len()
475 }
476
477 /// Whether the BVH was built (`with_bvh`).
478 pub fn has_bvh(&self) -> bool {
479 self.packed.has_bvh
480 }
481 }
482
483 /// A summary: the buffers are megabytes.
484 impl std::fmt::Debug for PreparedRtScene {
485 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
486 f.debug_struct("PreparedRtScene")
487 .field("triangles", &self.packed.tris.len())
488 .field("materials", &self.packed.materials.len())
489 .field("bvh_nodes", &self.packed.has_bvh.then_some(self.packed.nodes.len()))
490 .field("image", &self.image)
491 .finish()
492 }
493 }
494
495 /// The image a frame traces, as its parameter block describes it: the
496 /// texture's size and the quad's corners and opacity.
497 pub(crate) struct ParamImage {
498 pub width: u32,
499 pub height: u32,
500 pub corners: [[f32; 3]; 4],
501 pub opacity: f32,
502 }
503
504 fn v4([x, y, z]: [f32; 3]) -> [f32; 4] {
505 [x, y, z, 0.0]
506 }
507
508 /// One dispatch's parameter block. With no image (or one not uploaded
509 /// yet) the quad's opacity is 0, which lets every ray through it.
510 #[allow(clippy::too_many_arguments)]
511 pub(crate) fn rt_params(
512 camera: RtCamera,
513 size: (u32, u32),
514 sample_index: u32,
515 spp: u32,
516 image: Option<ParamImage>,
517 background: Option<[f32; 3]>,
518 environment: &RtEnvironment,
519 ) -> RtParams {
520 let (img_origin, img_u, img_v) = match image {
521 Some(ParamImage { width, height, corners: [tl, tr, _, bl], opacity }) => (
522 [tl[0], tl[1], tl[2], opacity.clamp(0.0, 1.0)],
523 [tr[0] - tl[0], tr[1] - tl[1], tr[2] - tl[2], width as f32],
524 [bl[0] - tl[0], bl[1] - tl[1], bl[2] - tl[2], height as f32],
525 ),
526 None => ([0.0; 4], [1.0, 0.0, 0.0, 1.0], [0.0, 1.0, 0.0, 1.0]),
527 };
528 RtParams {
529 inv_mvp: camera.inv_mvp,
530 width: size.0,
531 height: size.1,
532 sample_index,
533 max_bounces: MAX_BOUNCES,
534 spp,
535 _pad: [0; 3],
536 img_origin,
537 img_u,
538 img_v,
539 background: match background {
540 Some([r, g, b]) => [r, g, b, 1.0],
541 None => [0.0; 4],
542 },
543 sun_dir: v4(environment.sun_unit()),
544 sun_color: v4(environment.sun_color),
545 sky_zenith: v4(environment.sky_zenith),
546 sky_nadir: v4(environment.sky_nadir),
547 }
548 }
549
550 /// The denoiser's blocks, one per à-trous iteration (steps 1, 2, 4), for an
551 /// accumulation that will hold `n_after` samples once this frame's dispatch
552 /// lands — the colour sigma tightens as it grows.
553 pub(crate) fn denoise_params(width: u32, height: u32, n_after: u32) -> [DenoiseParams; DENOISE_ITERATIONS] {
554 let inv_sqrt_n = 1.0 / (n_after.max(1) as f32).sqrt();
555 std::array::from_fn(|i| DenoiseParams {
556 width,
557 height,
558 step: 1 << i,
559 first: (i == 0) as u32,
560 last: (i == DENOISE_ITERATIONS - 1) as u32,
561 inv_sqrt_n,
562 _pad: [0; 2],
563 })
564 }
565
566 #[cfg(test)]
567 mod tests {
568 use super::*;
569
570 // CPU mirror of the shader's traversal, for parity testing.
571 fn intersect_tri_cpu(ro: [f32; 3], rd: [f32; 3], t: &RtTriangle, t_limit: f32) -> f32 {
572 let sub = |a: [f32; 3], b: [f32; 3]| [a[0] - b[0], a[1] - b[1], a[2] - b[2]];
573 let cross = |a: [f32; 3], b: [f32; 3]| {
574 [
575 a[1] * b[2] - a[2] * b[1],
576 a[2] * b[0] - a[0] * b[2],
577 a[0] * b[1] - a[1] * b[0],
578 ]
579 };
580 let dot = |a: [f32; 3], b: [f32; 3]| a[0] * b[0] + a[1] * b[1] + a[2] * b[2];
581 let e1 = sub(t.p1, t.p0);
582 let e2 = sub(t.p2, t.p0);
583 let h = cross(rd, e2);
584 let a = dot(e1, h);
585 if a.abs() < 1e-8 {
586 return 1e30;
587 }
588 let f = 1.0 / a;
589 let s = sub(ro, t.p0);
590 let u = f * dot(s, h);
591 if !(0.0..=1.0).contains(&u) {
592 return 1e30;
593 }
594 let q = cross(s, e1);
595 let v = f * dot(rd, q);
596 if v < 0.0 || u + v > 1.0 {
597 return 1e30;
598 }
599 let tt = f * dot(e2, q);
600 if tt > 1e-4 && tt < t_limit {
601 return tt;
602 }
603 1e30
604 }
605
606 fn traverse_bvh_cpu(
607 nodes: &[GpuBvhNode],
608 tris: &[RtTriangle],
609 ro: [f32; 3],
610 rd: [f32; 3],
611 ) -> (f32, Option<usize>) {
612 if nodes.is_empty() {
613 return (1e30, None);
614 }
615 let inv = [1.0 / rd[0], 1.0 / rd[1], 1.0 / rd[2]];
616 let hit_aabb = |min: [f32; 3], max: [f32; 3], t_limit: f32| -> bool {
617 let mut tn = f32::NEG_INFINITY;
618 let mut tf = f32::INFINITY;
619 for a in 0..3 {
620 let t1 = (min[a] - ro[a]) * inv[a];
621 let t2 = (max[a] - ro[a]) * inv[a];
622 tn = tn.max(t1.min(t2));
623 tf = tf.min(t1.max(t2));
624 }
625 tf >= tn.max(0.0) && tn < t_limit
626 };
627 let mut best = 1e30f32;
628 let mut best_tri = None;
629 let mut stack = vec![0u32];
630 while let Some(idx) = stack.pop() {
631 let node = &nodes[idx as usize];
632 if !hit_aabb(node.min, node.max, best) {
633 continue;
634 }
635 if node.count > 0 {
636 for i in node.left_first..node.left_first + node.count {
637 let t = intersect_tri_cpu(ro, rd, &tris[i as usize], best);
638 if t < best {
639 best = t;
640 best_tri = Some(i as usize);
641 }
642 }
643 } else {
644 stack.push(node.left_first);
645 stack.push(node.left_first + 1);
646 }
647 }
648 (best, best_tri)
649 }
650
651 fn brute_force(tris: &[RtTriangle], ro: [f32; 3], rd: [f32; 3]) -> (f32, Option<usize>) {
652 let mut best = 1e30f32;
653 let mut best_tri = None;
654 for (i, t) in tris.iter().enumerate() {
655 let tt = intersect_tri_cpu(ro, rd, t, best);
656 if tt < best {
657 best = tt;
658 best_tri = Some(i);
659 }
660 }
661 (best, best_tri)
662 }
663
664 // Deterministic LCG so the test needs no rand dependency.
665 struct Lcg(u64);
666 impl Lcg {
667 fn next_f32(&mut self) -> f32 {
668 self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1442695040888963407);
669 ((self.0 >> 33) as f32) / (u32::MAX >> 1) as f32
670 }
671 fn point(&mut self, scale: f32) -> [f32; 3] {
672 [
673 (self.next_f32() - 0.5) * scale,
674 (self.next_f32() - 0.5) * scale,
675 (self.next_f32() - 0.5) * scale,
676 ]
677 }
678 }
679
680 fn random_scene(n: usize, seed: u64) -> Vec<RtTriangle> {
681 let mut rng = Lcg(seed);
682 (0..n)
683 .map(|i| {
684 let c = rng.point(20.0);
685 let jitter = |rng: &mut Lcg, c: [f32; 3]| {
686 let d = rng.point(2.0);
687 [c[0] + d[0], c[1] + d[1], c[2] + d[2]]
688 };
689 RtTriangle {
690 p0: jitter(&mut rng, c),
691 p1: jitter(&mut rng, c),
692 p2: jitter(&mut rng, c),
693 material: (i % 5) as u32,
694 }
695 })
696 .collect()
697 }
698
699 #[test]
700 fn test_bvh_matches_brute_force() {
701 let mut tris = random_scene(500, 42);
702 let nodes = build_bvh(&mut tris);
703 assert!(!nodes.is_empty());
704 let mut rng = Lcg(7);
705 let mut hits = 0;
706 for _ in 0..200 {
707 let ro = rng.point(40.0);
708 let target = rng.point(10.0);
709 let d = [target[0] - ro[0], target[1] - ro[1], target[2] - ro[2]];
710 let len = (d[0] * d[0] + d[1] * d[1] + d[2] * d[2]).sqrt().max(1e-6);
711 let rd = [d[0] / len, d[1] / len, d[2] / len];
712 let (t_bvh, tri_bvh) = traverse_bvh_cpu(&nodes, &tris, ro, rd);
713 let (t_ref, tri_ref) = brute_force(&tris, ro, rd);
714 assert_eq!(tri_bvh, tri_ref, "different triangle hit");
715 assert!((t_bvh - t_ref).abs() < 1e-4, "t mismatch: {t_bvh} vs {t_ref}");
716 if tri_bvh.is_some() {
717 hits += 1;
718 }
719 }
720 assert!(hits > 20, "test rays barely hit the scene ({hits}/200)");
721 }
722
723 #[test]
724 fn test_bvh_leaf_ranges_cover_all_triangles() {
725 let mut tris = random_scene(300, 9);
726 let nodes = build_bvh(&mut tris);
727 let mut seen = vec![false; tris.len()];
728 for node in &nodes {
729 if node.count > 0 {
730 for i in node.left_first..node.left_first + node.count {
731 assert!(!seen[i as usize], "triangle {i} in two leaves");
732 seen[i as usize] = true;
733 }
734 }
735 }
736 assert!(seen.iter().all(|&s| s), "not every triangle is in a leaf");
737 }
738
739 #[test]
740 fn test_bvh_degenerate_identical_centroids() {
741 // All triangles share one centroid: SAH can't split, the median
742 // fallback must still terminate and cover everything.
743 let tri = RtTriangle {
744 p0: [0.0, 0.0, 0.0],
745 p1: [1.0, 0.0, 0.0],
746 p2: [0.0, 1.0, 0.0],
747 material: 0,
748 };
749 let mut tris = vec![tri; 100];
750 let nodes = build_bvh(&mut tris);
751 let covered: u32 = nodes.iter().filter(|n| n.count > 0).map(|n| n.count).sum();
752 assert_eq!(covered, 100);
753 let (t, hit) = traverse_bvh_cpu(&nodes, &tris, [0.2, 0.2, -5.0], [0.0, 0.0, 1.0]);
754 assert!(hit.is_some());
755 assert!((t - 5.0).abs() < 1e-3);
756 }
757
758 #[test]
759 fn test_bvh_empty_and_single() {
760 let mut empty: Vec<RtTriangle> = Vec::new();
761 assert!(build_bvh(&mut empty).is_empty());
762
763 let mut single = vec![RtTriangle {
764 p0: [-1.0, -1.0, 0.0],
765 p1: [1.0, -1.0, 0.0],
766 p2: [0.0, 1.0, 0.0],
767 material: 3,
768 }];
769 let nodes = build_bvh(&mut single);
770 assert_eq!(nodes.len(), 1);
771 assert_eq!(nodes[0].count, 1);
772 let (t, hit) = traverse_bvh_cpu(&nodes, &single, [0.0, 0.0, -3.0], [0.0, 0.0, 1.0]);
773 assert_eq!(hit, Some(0));
774 assert!((t - 3.0).abs() < 1e-4);
775 }
776
777 #[test]
778 fn test_prepared_scene_bvh_built_late_matches_built_early() {
779 // A scene prepared for the ray-query tier and uploaded to the
780 // compute one gets the BVH it would have been prepared with.
781 let tris = random_scene(400, 3);
782 let mats = [RtMaterial { albedo: [0.5; 3], emission: [0.0; 3] }; 5];
783 let image = RtImage {
784 image: 1,
785 corners: [[-1.0, 1.0, 0.0], [1.0, 1.0, 0.0], [1.0, -1.0, 0.0], [-1.0, -1.0, 0.0]],
786 opacity: 1.0,
787 };
788 let early = PreparedRtScene::new(tris.clone(), &mats, Some(image), true);
789 let late = PreparedRtScene::new(tris, &mats, Some(image), false);
790 assert!(early.has_bvh() && !late.has_bvh());
791 assert!(late.packed.nodes.is_empty());
792 assert_eq!(early.triangle_count(), 402, "the image's quad joins the scene");
793 let built = late.packed.with_bvh();
794 assert!(built.has_bvh);
795 let bytes = |s: &PackedScene| {
796 (bytemuck::cast_slice::<_, u8>(&s.tris).to_vec(), bytemuck::cast_slice::<_, u8>(&s.nodes).to_vec())
797 };
798 assert_eq!(bytes(&built), bytes(&early.packed));
799 assert!(matches!(early.packed.with_bvh(), std::borrow::Cow::Borrowed(_)));
800 }
801
802 #[test]
803 fn test_prepared_scene_crosses_threads() {
804 fn send_sync<T: Send + Sync>() {}
805 send_sync::<PreparedRtScene>();
806 let prepared = std::thread::spawn(|| PreparedRtScene::new(random_scene(50, 1), &[], None, true))
807 .join()
808 .unwrap();
809 assert_eq!(prepared.triangle_count(), 50);
810 }
811 }