diff --git a/README.md b/README.md index 3d034ee..7bedbab 100644 --- a/README.md +++ b/README.md @@ -93,13 +93,13 @@ It uses all cores via rayon, shows a progress bar and (for the `cornell` scene) | `spheres` | Large scene: ~470 random diffuse/mirror/glass spheres under a sky (after *Ray Tracing in One Weekend*). The BVH makes it ~11x faster than testing every object | | `night` | The same field at night, lit by 61 glowing spheres. Each shading point picks one light, weighted by power / distance², so many lights cost no more shadow rays than one | | `classic` | The classic Cornell box: one-sided ceiling panel light, two boxes, colour bleeding from the walls | -| `city` | Massive scale: a procedural 400x400-lot city of ~124,000 objects (buildings with stepped tiers, mirror-glass towers, parks with trees) under a low sun. Renders 800x600 at 64 spp in ~2.5 s; without the BVH it is ~1000x slower | +| `city` | Massive scale: a procedural 400x400-lot city of ~124,000 objects (buildings with stepped tiers, mirror-glass towers, parks with trees) under a low sun. Renders 800x600 at 64 spp in ~1 s; without the BVH it is ~1000x slower | Scenes are defined in [`rust/src/scenes.rs`](rust/src/scenes.rs) with a small builder (`plane`, `sphere`, `cuboid`, `box_light`, `sphere_light`, `rect_light`). Every emitter must be added with a `*_light` method so light sampling knows about it. ### Beyond the Python version -- **BVH** - bounded objects live in a bounding volume hierarchy (median split); planes are tested separately. +- **8-wide BVH with AVX-512** - bounded objects live in a bounding volume hierarchy (median split, collapsed from binary to 8 children per node); planes are tested separately. Each node stores its children's boxes as structure-of-arrays, so one AVX-512 sequence tests a ray against all 8. Leaves hold up to 8 spheres, also intersected 8 at a time, but at most one other shape (boxes and rects are still tested one by one, so tight leaves are faster for them). The CPU is checked at runtime; without AVX-512, or with `PATHTRACER_NO_AVX512=1` set, the same traversal runs with scalar kernels and renders bit-identical images. Compared with the earlier binary BVH on a Ryzen 7 9800X3D: `city` 2.6x, `spheres` 2.1x, `night` 1.8x faster; the small scenes are unchanged within ±10%. - **Many-light sampling** - box, sphere and one-sided rectangle lights, one picked per shading point in proportion to its estimated contribution. - **Russian roulette** - after 3 bounces, dim paths are terminated randomly (survivors are re-weighted, so the image stays unbiased). About 2.5x more image quality per second of rendering in this scene. diff --git a/rust/src/bvh.rs b/rust/src/bvh.rs index ec1d6bc..298fa95 100644 --- a/rust/src/bvh.rs +++ b/rust/src/bvh.rs @@ -1,5 +1,15 @@ //! Bounding volume hierarchy over the scene's bounded objects, so a ray only //! tests the objects whose boxes it passes through instead of all of them. +//! +//! The tree is built as a binary BVH (median split) and then collapsed into +//! an 8-wide BVH: each node holds the bounds of up to 8 children in +//! structure-of-arrays form, so one AVX-512 instruction sequence tests a ray +//! against all 8 boxes. Leaves hold up to 8 objects; spheres in a leaf are +//! intersected 8 at a time the same way, other shapes one by one through a +//! callback. Without AVX-512 (or with PATHTRACER_NO_AVX512 set) the same +//! traversal runs with scalar kernels. + +use std::mem::MaybeUninit; use crate::scene::Ray; use crate::vec3::Vec3; @@ -27,159 +37,651 @@ impl Aabb { (self.min + self.max) * 0.5 } - /// Slab test. `inv_dir` components may be infinite for axis-parallel rays; - /// the resulting NaNs are discarded by f64::min/max. - fn hit(&self, origin: Vec3, inv_dir: Vec3, mut t_min: f64, mut t_max: f64) -> bool { - for axis in 0..3 { - let mut t0 = (self.min[axis] - origin[axis]) * inv_dir[axis]; - let mut t1 = (self.max[axis] - origin[axis]) * inv_dir[axis]; - if inv_dir[axis] < 0.0 { - std::mem::swap(&mut t0, &mut t1); - } - t_min = t_min.max(t0); - t_max = t_max.min(t1); - if t_max < t_min { - return false; - } - } - true + fn surface_area(&self) -> f64 { + let d = self.max - self.min; + 2.0 * (d.x * d.y + d.y * d.z + d.x * d.z) } } -const LEAF_SIZE: usize = 4; +const WIDTH: usize = 8; +/// Leaves hold up to WIDTH objects, since spheres are tested WIDTH at a time, +/// but at most this many other shapes, which are tested one by one: for +/// box-heavy scenes, small leaves (more SIMD node tests, fewer scalar box +/// tests) measured fastest. +const MAX_OTHERS_PER_LEAF: usize = 1; -/// Flattened node: a leaf covers `prims[start..start + count]`; an interior -/// node (count == 0) has its left child at the next index and its right child at `right`. -struct Node { +// --------------------------------------------------------------------------- +// Build: binary BVH first. + +struct BinNode { bounds: Aabb, - start: u32, - count: u32, - right: u32, - axis: u8, + start: usize, + count: usize, // > 0 for leaves + left: usize, + right: usize, +} + +/// Per-object inputs to the build. +struct BuildInput<'a> { + bounds: &'a [Aabb], + centroids: Vec, + is_sphere: Vec, +} + +fn build_binary(nodes: &mut Vec, prims: &mut [u32], offset: usize, input: &BuildInput) -> usize { + let (bounds, centroids) = (input.bounds, &input.centroids); + let node_bounds = prims.iter().fold(Aabb::EMPTY, |b, &p| b.union(bounds[p as usize])); + let idx = nodes.len(); + nodes.push(BinNode { bounds: node_bounds, start: offset, count: prims.len(), left: 0, right: 0 }); + let others = prims.iter().filter(|&&p| !input.is_sphere[p as usize]).count(); + if prims.len() <= WIDTH && others <= MAX_OTHERS_PER_LEAF { + return idx; + } + + // Median split along the axis where the object centroids spread widest. + let centroid_bounds = prims.iter().fold(Aabb::EMPTY, |b, &p| { + let c = centroids[p as usize]; + b.union(Aabb { min: c, max: c }) + }); + let extent = centroid_bounds.max - centroid_bounds.min; + let axis = if extent.x >= extent.y && extent.x >= extent.z { + 0 + } else if extent.y >= extent.z { + 1 + } else { + 2 + }; + let half = prims.len() / 2; + // If all centroids coincide the order doesn't matter; just halve. + if extent[axis] > 0.0 { + prims.select_nth_unstable_by(half, |&a, &b| centroids[a as usize][axis].total_cmp(¢roids[b as usize][axis])); + } + let (lo, hi) = prims.split_at_mut(half); + let left = build_binary(nodes, lo, offset, input); + let right = build_binary(nodes, hi, offset + half, input); + nodes[idx].count = 0; + nodes[idx].left = left; + nodes[idx].right = right; + idx +} + +// --------------------------------------------------------------------------- +// 8-wide nodes. + +/// Child references: an interior child is its node index; a leaf child has +/// LEAF_BIT set, flags saying whether it holds spheres and/or other shapes +/// (so traversal can skip the pass it doesn't need), (count - 1) in bits +/// 26..29 and its first slot in bits 0..26. +const LEAF_BIT: u32 = 1 << 31; +const HAS_OTHERS: u32 = 1 << 30; +const HAS_SPHERES: u32 = 1 << 29; +const START_BITS: u32 = 26; +const START_MASK: u32 = (1 << START_BITS) - 1; + +fn leaf_ref(start: usize, count: usize, has_spheres: bool, has_others: bool) -> u32 { + assert!(start < (1 << START_BITS), "too many objects for the BVH leaf encoding"); + debug_assert!((1..=WIDTH).contains(&count)); + let mut r = LEAF_BIT | (((count - 1) as u32) << START_BITS) | start as u32; + if has_spheres { + r |= HAS_SPHERES; + } + if has_others { + r |= HAS_OTHERS; + } + r +} + +fn decode_leaf(r: u32) -> (usize, usize) { + ((r & START_MASK) as usize, (((r >> START_BITS) & 7) + 1) as usize) +} + +/// Child bounds in structure-of-arrays layout. Unused lanes have all bounds +/// at +inf: every slab distance is then ±inf, so the lane can never be hit +/// and needs no separate mask, provided t_max is finite (with t_max = inf a +/// ray with all-positive direction would see inf <= inf and "hit" it, which +/// is why the traversal clamps t_max to f64::MAX). +#[repr(C, align(64))] +struct Node8 { + min_x: [f64; WIDTH], + min_y: [f64; WIDTH], + min_z: [f64; WIDTH], + max_x: [f64; WIDTH], + max_y: [f64; WIDTH], + max_z: [f64; WIDTH], + child: [u32; WIDTH], +} + +impl Node8 { + fn empty() -> Self { + let inf = [f64::INFINITY; WIDTH]; + Node8 { min_x: inf, min_y: inf, min_z: inf, max_x: inf, max_y: inf, max_z: inf, child: [0; WIDTH] } + } +} + +/// Sphere data per slot, in BVH order, padded by WIDTH entries so a leaf's +/// 8 lanes can always be loaded. Non-sphere slots have r2 = -inf, which +/// makes the discriminant -inf so they never register a hit. +struct Spheres { + cx: Vec, + cy: Vec, + cz: Vec, + r2: Vec, +} + +/// Per-ray values shared by all box and sphere tests. +#[derive(Clone, Copy)] +struct RayData { + o: Vec3, + d: Vec3, + inv: Vec3, + /// d·d, as the scalar sphere test computes it. + a: f64, +} + +impl RayData { + fn new(r: &Ray) -> Self { + // A finite stand-in for 1/0 keeps (bound - origin) * inv free of the + // NaN that 0 * inf would produce. + let inv = |d: f64| { + let i = 1.0 / d; + if i.is_finite() { + i + } else { + 1e300_f64.copysign(i) + } + }; + RayData { + o: r.origin, + d: r.direction, + inv: Vec3::new(inv(r.direction.x), inv(r.direction.y), inv(r.direction.z)), + a: r.direction.dot(r.direction), + } + } +} + +// --------------------------------------------------------------------------- +// Scalar kernels. + +/// Returns a bitmask of the child boxes the ray hits within [t_min, t_max], +/// writing each child's entry distance to `tnear`. +#[inline(always)] +fn node_scalar(n: &Node8, r: &RayData, t_min: f64, t_max: f64, tnear: &mut [f64; WIDTH]) -> u8 { + let mut mask = 0u8; + for i in 0..WIDTH { + let (x0, x1) = ((n.min_x[i] - r.o.x) * r.inv.x, (n.max_x[i] - r.o.x) * r.inv.x); + let (y0, y1) = ((n.min_y[i] - r.o.y) * r.inv.y, (n.max_y[i] - r.o.y) * r.inv.y); + let (z0, z1) = ((n.min_z[i] - r.o.z) * r.inv.z, (n.max_z[i] - r.o.z) * r.inv.z); + let near = x0.min(x1).max(y0.min(y1)).max(z0.min(z1).max(t_min)); + let far = x0.max(x1).min(y0.max(y1)).min(z0.max(z1).min(t_max)); + tnear[i] = near; + if near <= far { + mask |= 1 << i; + } + } + mask +} + +/// Intersects the ray with the spheres in slots start..start + count. +/// Returns a bitmask of lanes hit within [t_min, t_max] and each lane's +/// hit distance. Mirrors Shape::hit for spheres operation by operation. +#[inline(always)] +fn spheres_scalar(s: &Spheres, start: usize, count: usize, r: &RayData, t_min: f64, t_max: f64) -> (u8, [f64; WIDTH]) { + let mut mask = 0u8; + let mut roots = [0.0; WIDTH]; + for i in 0..count { + let k = start + i; + let oc = Vec3::new(r.o.x - s.cx[k], r.o.y - s.cy[k], r.o.z - s.cz[k]); + let b = oc.dot(r.d); + let c = oc.dot(oc) - s.r2[k]; + let disc = b * b - r.a * c; + if disc < 0.0 || disc.is_nan() { + continue; + } + let sq = disc.sqrt(); + let mut root = (-b - sq) / r.a; + if root < t_min || root > t_max { + root = (-b + sq) / r.a; + if root < t_min || root > t_max { + continue; + } + } + roots[i] = root; + mask |= 1 << i; + } + (mask, roots) +} + +// --------------------------------------------------------------------------- +// AVX-512 kernels: the same arithmetic on 8 lanes of f64 at once. + +#[cfg(target_arch = "x86_64")] +mod avx512 { + use super::{Node8, RayData, Spheres, WIDTH}; + use std::arch::x86_64::*; + + #[target_feature(enable = "avx512f")] + #[inline] + pub unsafe fn node(n: &Node8, r: &RayData, t_min: f64, t_max: f64, tnear: &mut [f64; WIDTH]) -> u8 { + let slab = |lo: &[f64; WIDTH], hi: &[f64; WIDTH], o: f64, inv: f64| { + let (o, inv) = (_mm512_set1_pd(o), _mm512_set1_pd(inv)); + let t0 = _mm512_mul_pd(_mm512_sub_pd(_mm512_loadu_pd(lo.as_ptr()), o), inv); + let t1 = _mm512_mul_pd(_mm512_sub_pd(_mm512_loadu_pd(hi.as_ptr()), o), inv); + (_mm512_min_pd(t0, t1), _mm512_max_pd(t0, t1)) + }; + let (nx, fx) = slab(&n.min_x, &n.max_x, r.o.x, r.inv.x); + let (ny, fy) = slab(&n.min_y, &n.max_y, r.o.y, r.inv.y); + let (nz, fz) = slab(&n.min_z, &n.max_z, r.o.z, r.inv.z); + let near = _mm512_max_pd(_mm512_max_pd(nx, ny), _mm512_max_pd(nz, _mm512_set1_pd(t_min))); + let far = _mm512_min_pd(_mm512_min_pd(fx, fy), _mm512_min_pd(fz, _mm512_set1_pd(t_max))); + _mm512_storeu_pd(tnear.as_mut_ptr(), near); + _mm512_cmp_pd_mask::<_CMP_LE_OQ>(near, far) + } + + #[target_feature(enable = "avx512f")] + #[inline] + pub unsafe fn spheres( + s: &Spheres, + start: usize, + count: usize, + r: &RayData, + t_min: f64, + t_max: f64, + ) -> (u8, [f64; WIDTH]) { + let mut roots = [0.0; WIDTH]; + let lanes = ((1u32 << count) - 1) as u8; + let load = |v: &Vec| _mm512_loadu_pd(v.as_ptr().add(start)); + let (add, sub, mul) = ( + |a, b| _mm512_add_pd(a, b), + |a, b| _mm512_sub_pd(a, b), + |a, b| _mm512_mul_pd(a, b), + ); + + let ocx = sub(_mm512_set1_pd(r.o.x), load(&s.cx)); + let ocy = sub(_mm512_set1_pd(r.o.y), load(&s.cy)); + let ocz = sub(_mm512_set1_pd(r.o.z), load(&s.cz)); + let (dx, dy, dz) = (_mm512_set1_pd(r.d.x), _mm512_set1_pd(r.d.y), _mm512_set1_pd(r.d.z)); + let b = add(add(mul(ocx, dx), mul(ocy, dy)), mul(ocz, dz)); + let c = sub(add(add(mul(ocx, ocx), mul(ocy, ocy)), mul(ocz, ocz)), load(&s.r2)); + let a = _mm512_set1_pd(r.a); + let disc = sub(mul(b, b), mul(a, c)); + let zero = _mm512_setzero_pd(); + let mut mask = lanes & _mm512_cmp_pd_mask::<_CMP_GE_OQ>(disc, zero); + if mask == 0 { + return (0, roots); + } + + let sq = _mm512_sqrt_pd(disc); + let neg_b = sub(zero, b); + let r1 = _mm512_div_pd(sub(neg_b, sq), a); + let r2 = _mm512_div_pd(add(neg_b, sq), a); + let (lo, hi) = (_mm512_set1_pd(t_min), _mm512_set1_pd(t_max)); + let in_range = |t| _mm512_cmp_pd_mask::<_CMP_GE_OQ>(t, lo) & _mm512_cmp_pd_mask::<_CMP_LE_OQ>(t, hi); + let ok1 = in_range(r1); + let ok2 = in_range(r2); + // Nearer root when it is in range, as in the scalar version. + let root = _mm512_mask_blend_pd(ok1, r2, r1); + mask &= ok1 | ok2; + _mm512_storeu_pd(roots.as_mut_ptr(), root); + (mask, roots) + } +} + +// --------------------------------------------------------------------------- +// Traversal, written once and instantiated for each set of kernels (a macro +// rather than generics so the AVX-512 kernels inline into a function that is +// itself compiled with AVX-512 enabled). + +/// Traversal pushes at most WIDTH - 1 entries per level beyond the one it pops, +/// so this covers trees up to (STACK_SIZE - 1) / 7 = 36 levels deep (checked in build). +const STACK_SIZE: usize = 256; + +macro_rules! traversal { + ($(#[$attr:meta])* $closest:ident, $any:ident, $node:path, $spheres:path) => { + /// Closest hit. `other(object, t_max)` intersects non-sphere objects + /// and returns the distance if hit closer than t_max. Returns + /// (object, t, whether it was a sphere). + $(#[$attr])* + unsafe fn $closest( + &self, + root: u32, + r: &RayData, + t_min: f64, + mut t_max: f64, + mut other: impl FnMut(usize, f64) -> Option, + ) -> Option<(usize, f64, bool)> { + let mut best = None; + // Uninitialised: zero-filling these on every ray is measurable in + // small scenes. Only entries below sp are ever read. + let mut stack_ref = [MaybeUninit::::uninit(); STACK_SIZE]; + let mut stack_t = [MaybeUninit::::uninit(); STACK_SIZE]; + stack_ref[0].write(root); + stack_t[0].write(t_min); + let mut sp = 1; + let mut tnear = [0.0; WIDTH]; + + while sp > 0 { + sp -= 1; + if stack_t[sp].assume_init() > t_max { + continue; // Already found something closer than this whole subtree. + } + let node_ref = stack_ref[sp].assume_init(); + if node_ref & LEAF_BIT != 0 { + let (start, count) = decode_leaf(node_ref); + if node_ref & HAS_SPHERES != 0 { + let (mut mask, roots) = $spheres(&self.spheres, start, count, r, t_min, t_max); + while mask != 0 { + let i = mask.trailing_zeros() as usize; + mask &= mask - 1; + if roots[i] < t_max { + t_max = roots[i]; + best = Some((self.slots[start + i] as usize, roots[i], true)); + } + } + } + if node_ref & HAS_OTHERS != 0 { + for k in start..start + count { + if self.spheres.r2[k] == f64::NEG_INFINITY { + let obj = self.slots[k] as usize; + if let Some(t) = other(obj, t_max) { + t_max = t; + best = Some((obj, t, false)); + } + } + } + } + continue; + } + + let node = &self.nodes[node_ref as usize]; + let mut mask = $node(node, r, t_min, t_max, &mut tnear); + // Push hit children farthest first, so the nearest is popped + // next and shrinks t_max soonest. + let mut hits = [(0u32, 0.0f64); WIDTH]; + let mut n = 0; + while mask != 0 { + let i = mask.trailing_zeros() as usize; + mask &= mask - 1; + let mut j = n; + while j > 0 && hits[j - 1].1 < tnear[i] { + hits[j] = hits[j - 1]; + j -= 1; + } + hits[j] = (node.child[i], tnear[i]); + n += 1; + } + debug_assert!(sp + n <= STACK_SIZE); + for &(child, t) in &hits[..n] { + stack_ref[sp].write(child); + stack_t[sp].write(t); + sp += 1; + } + } + best + } + + /// True as soon as anything is hit within (t_min, t_max); + /// `other(object)` tests non-sphere objects. + $(#[$attr])* + unsafe fn $any( + &self, + root: u32, + r: &RayData, + t_min: f64, + t_max: f64, + mut other: impl FnMut(usize) -> bool, + ) -> bool { + let mut stack = [MaybeUninit::::uninit(); STACK_SIZE]; + stack[0].write(root); + let mut sp = 1; + let mut tnear = [0.0; WIDTH]; + + while sp > 0 { + sp -= 1; + let node_ref = stack[sp].assume_init(); + if node_ref & LEAF_BIT != 0 { + let (start, count) = decode_leaf(node_ref); + if node_ref & HAS_SPHERES != 0 && $spheres(&self.spheres, start, count, r, t_min, t_max).0 != 0 { + return true; + } + if node_ref & HAS_OTHERS != 0 { + for k in start..start + count { + if self.spheres.r2[k] == f64::NEG_INFINITY && other(self.slots[k] as usize) { + return true; + } + } + } + continue; + } + + let node = &self.nodes[node_ref as usize]; + let mut mask = $node(node, r, t_min, t_max, &mut tnear); + while mask != 0 { + let i = mask.trailing_zeros() as usize; + mask &= mask - 1; + debug_assert!(sp < STACK_SIZE); + stack[sp].write(node.child[i]); + sp += 1; + } + } + false + } + }; } pub struct Bvh { - nodes: Vec, - /// Object indices, reordered so each leaf's objects are contiguous. - prims: Vec, + /// Reference to the root: node 0, or a single leaf when the whole scene + /// fits in one (then there is no node test at all). None if empty. + root: Option, + /// Bounds of everything; tested first when the root is a single leaf, + /// which otherwise has no box test to reject rays that miss it. + root_bounds: Aabb, + nodes: Vec, + /// Object index for each slot; each leaf covers a contiguous slot range. + slots: Vec, + spheres: Spheres, + use_avx512: bool, } impl Bvh { - pub fn build(bounds: &[Aabb]) -> Bvh { - let mut bvh = Bvh { - nodes: Vec::with_capacity(2 * bounds.len()), - prims: (0..bounds.len() as u32).collect(), - }; + /// `spheres[i]` is Some((center, radius)) if object i is a sphere. + pub fn build(bounds: &[Aabb], spheres: &[Option<(Vec3, f64)>]) -> Bvh { + let mut slots: Vec = (0..bounds.len() as u32).collect(); + let mut bin = Vec::with_capacity(2 * bounds.len()); if !bounds.is_empty() { - let centroids: Vec = bounds.iter().map(Aabb::centroid).collect(); - bvh.build_range(0, bounds.len(), bounds, ¢roids); + let input = BuildInput { + bounds, + centroids: bounds.iter().map(Aabb::centroid).collect(), + is_sphere: spheres.iter().map(Option::is_some).collect(), + }; + build_binary(&mut bin, &mut slots, 0, &input); } + + let mut soa = Spheres { + cx: vec![0.0; slots.len() + WIDTH], + cy: vec![0.0; slots.len() + WIDTH], + cz: vec![0.0; slots.len() + WIDTH], + r2: vec![f64::NEG_INFINITY; slots.len() + WIDTH], + }; + for (k, &obj) in slots.iter().enumerate() { + if let Some((c, radius)) = spheres[obj as usize] { + (soa.cx[k], soa.cy[k], soa.cz[k], soa.r2[k]) = (c.x, c.y, c.z, radius * radius); + } + } + + let mut bvh = Bvh { + root: None, + root_bounds: bin.first().map_or(Aabb::EMPTY, |b| b.bounds), + nodes: Vec::new(), + slots, + spheres: soa, + use_avx512: detect_avx512(), + }; + if let Some(b) = bin.first() { + bvh.root = Some(if b.count > 0 { bvh.leaf(b) } else { bvh.collapse(&bin, 0) as u32 }); + } + assert!( + (WIDTH - 1) * bvh.depth() + 1 <= STACK_SIZE, + "BVH too deep for the traversal stack" + ); bvh } - /// Median split along the axis where the object centroids spread widest. - fn build_range(&mut self, start: usize, end: usize, bounds: &[Aabb], centroids: &[Vec3]) -> usize { - let prims = &self.prims[start..end]; - let node_bounds = prims.iter().fold(Aabb::EMPTY, |b, &p| b.union(bounds[p as usize])); - let centroid_bounds = prims.iter().fold(Aabb::EMPTY, |b, &p| { - let c = centroids[p as usize]; - b.union(Aabb { min: c, max: c }) - }); - let extent = centroid_bounds.max - centroid_bounds.min; - let axis = if extent.x >= extent.y && extent.x >= extent.z { - 0 - } else if extent.y >= extent.z { - 1 - } else { - 2 - }; + fn leaf(&self, b: &BinNode) -> u32 { + let r2 = &self.spheres.r2[b.start..b.start + b.count]; + let has_spheres = r2.iter().any(|&r| r != f64::NEG_INFINITY); + let has_others = r2.iter().any(|&r| r == f64::NEG_INFINITY); + leaf_ref(b.start, b.count, has_spheres, has_others) + } - let idx = self.nodes.len(); - self.nodes.push(Node { - bounds: node_bounds, - start: start as u32, - count: (end - start) as u32, - right: 0, - axis: axis as u8, - }); - if end - start <= LEAF_SIZE || extent[axis] <= 0.0 { - return idx; + /// Converts binary node `b` (and its subtree) into 8-wide nodes, returning + /// the new node's index. Starts from b's two children and repeatedly + /// opens up the interior child with the largest surface area until there + /// are 8 children or only leaves remain. + fn collapse(&mut self, bin: &[BinNode], b: usize) -> usize { + let mut children = vec![bin[b].left, bin[b].right]; + while children.len() < WIDTH { + let Some(pos) = children + .iter() + .enumerate() + .filter(|(_, &c)| bin[c].count == 0) + .max_by(|(_, &a), (_, &c)| bin[a].bounds.surface_area().total_cmp(&bin[c].bounds.surface_area())) + .map(|(pos, _)| pos) + else { + break; + }; + let opened = children.swap_remove(pos); + children.push(bin[opened].left); + children.push(bin[opened].right); } - let half = (end - start) / 2; - self.prims[start..end].select_nth_unstable_by(half, |&a, &b| { - centroids[a as usize][axis].total_cmp(¢roids[b as usize][axis]) - }); - self.build_range(start, start + half, bounds, centroids); - let right = self.build_range(start + half, end, bounds, centroids); - self.nodes[idx].count = 0; - self.nodes[idx].right = right as u32; + let idx = self.nodes.len(); + self.nodes.push(Node8::empty()); + for (lane, &c) in children.iter().enumerate() { + let child_ref = if bin[c].count > 0 { + self.leaf(&bin[c]) + } else { + self.collapse(bin, c) as u32 + }; + let (lo, hi) = (bin[c].bounds.min, bin[c].bounds.max); + let n = &mut self.nodes[idx]; + (n.min_x[lane], n.min_y[lane], n.min_z[lane]) = (lo.x, lo.y, lo.z); + (n.max_x[lane], n.max_y[lane], n.max_z[lane]) = (hi.x, hi.y, hi.z); + n.child[lane] = child_ref; + } idx } - /// Calls `visit` for each object in a leaf whose box the ray passes through - /// within (t_min, t_max). `visit` returns Some(t) to shrink t_max after a - /// hit, or Err(()) to stop traversal entirely. - fn traverse( + /// Number of 8-wide node levels (0 when the root is a single leaf). + pub fn depth(&self) -> usize { + fn go(b: &Bvh, n: usize) -> usize { + 1 + b.nodes[n] + .child + .iter() + .zip(&b.nodes[n].min_x) + .filter(|(&c, &mx)| mx != f64::INFINITY && c & LEAF_BIT == 0) + .map(|(&c, _)| go(b, c as usize)) + .max() + .unwrap_or(0) + } + if self.nodes.is_empty() { + 0 + } else { + go(self, 0) + } + } + + /// False if the root is a single leaf whose bounds the ray misses. (A + /// node root does this test for its children itself.) + #[inline(always)] + fn root_hit(&self, root: u32, r: &RayData, t_min: f64, t_max: f64) -> bool { + if root & LEAF_BIT == 0 { + return true; + } + let (lo, hi) = (self.root_bounds.min, self.root_bounds.max); + let (mut near, mut far) = (t_min, t_max); + for axis in 0..3 { + let t0 = (lo[axis] - r.o[axis]) * r.inv[axis]; + let t1 = (hi[axis] - r.o[axis]) * r.inv[axis]; + near = near.max(t0.min(t1)); + far = far.min(t0.max(t1)); + } + near <= far + } + + /// Whether traversal runs the AVX-512 kernels. When the whole scene is a + /// single leaf without spheres there is nothing for them to do, and code + /// compiled for AVX-512 measured slightly slower on such scenes. + fn simd_pays_off(&self, root: u32) -> bool { + self.use_avx512 && (root & LEAF_BIT == 0 || root & HAS_SPHERES != 0) + } + + pub fn uses_avx512(&self) -> bool { + self.root.is_some_and(|root| self.simd_pays_off(root)) + } + + traversal!(closest_scalar, any_scalar, node_scalar, spheres_scalar); + + #[cfg(target_arch = "x86_64")] + traversal!( + #[target_feature(enable = "avx512f")] + closest_avx512, + any_avx512, + avx512::node, + avx512::spheres + ); + + /// See the traversal macro. Returns (object, t, is_sphere) for the closest hit. + pub fn closest( &self, r: &Ray, t_min: f64, - mut t_max: f64, - mut visit: impl FnMut(usize, f64) -> Result, ()>, - ) { - if self.nodes.is_empty() { - return; + t_max: f64, + other: impl FnMut(usize, f64) -> Option, + ) -> Option<(usize, f64, bool)> { + let Some(root) = self.root else { + return None; + }; + let rd = RayData::new(r); + let t_max = t_max.min(f64::MAX); // See Node8. + if !self.root_hit(root, &rd, t_min, t_max) { + return None; } - let inv_dir = Vec3::new(1.0 / r.direction.x, 1.0 / r.direction.y, 1.0 / r.direction.z); - let mut stack = [0u32; 64]; - let mut sp = 1; - while sp > 0 { - sp -= 1; - let node = &self.nodes[stack[sp] as usize]; - if !node.bounds.hit(r.origin, inv_dir, t_min, t_max) { - continue; - } - if node.count > 0 { - for &p in &self.prims[node.start as usize..(node.start + node.count) as usize] { - match visit(p as usize, t_max) { - Ok(Some(t)) => t_max = t, - Ok(None) => {} - Err(()) => return, - } - } - } else { - // Push the far child first so the near one is visited first, - // which shrinks t_max sooner. - let left = stack[sp] + 1; - let (near, far) = if r.direction[node.axis as usize] < 0.0 { - (node.right, left) - } else { - (left, node.right) - }; - stack[sp] = far; - stack[sp + 1] = near; - sp += 2; - } + #[cfg(target_arch = "x86_64")] + if self.simd_pays_off(root) { + // SAFETY: simd_pays_off() requires use_avx512, which is only set when the CPU supports AVX-512F. + return unsafe { self.closest_avx512(root, &rd, t_min, t_max, other) }; } + unsafe { self.closest_scalar(root, &rd, t_min, t_max, other) } } - /// Closest hit: `hit(object, t_max)` returns the hit distance if the object - /// is hit closer than t_max. - pub fn closest(&self, r: &Ray, t_min: f64, t_max: f64, mut hit: impl FnMut(usize, f64) -> Option) { - self.traverse(r, t_min, t_max, |p, t_max| Ok(hit(p, t_max))); - } - - /// True as soon as `hit(object)` reports any hit. - pub fn any(&self, r: &Ray, t_min: f64, t_max: f64, mut hit: impl FnMut(usize) -> bool) -> bool { - let mut found = false; - self.traverse(r, t_min, t_max, |p, _| { - if hit(p) { - found = true; - Err(()) - } else { - Ok(None) - } - }); - found + pub fn any(&self, r: &Ray, t_min: f64, t_max: f64, other: impl FnMut(usize) -> bool) -> bool { + let Some(root) = self.root else { + return false; + }; + let rd = RayData::new(r); + let t_max = t_max.min(f64::MAX); // See Node8. + if !self.root_hit(root, &rd, t_min, t_max) { + return false; + } + #[cfg(target_arch = "x86_64")] + if self.simd_pays_off(root) { + // SAFETY: as in closest(). + return unsafe { self.any_avx512(root, &rd, t_min, t_max, other) }; + } + unsafe { self.any_scalar(root, &rd, t_min, t_max, other) } + } +} + +fn detect_avx512() -> bool { + if std::env::var_os("PATHTRACER_NO_AVX512").is_some() { + return false; + } + #[cfg(target_arch = "x86_64")] + { + std::arch::is_x86_feature_detected!("avx512f") + } + #[cfg(not(target_arch = "x86_64"))] + { + false } } diff --git a/rust/src/main.rs b/rust/src/main.rs index bb70b6f..3b2a85d 100644 --- a/rust/src/main.rs +++ b/rust/src/main.rs @@ -409,6 +409,11 @@ fn main() { println!("Adaptive sampling: off"); } println!("Threads: {}", rayon::current_num_threads()); + println!( + "BVH: 8-wide, depth {}, {}", + scene.bvh_depth(), + if scene.uses_avx512() { "AVX-512" } else { "scalar" } + ); println!("Output: {}", settings.output); println!(); diff --git a/rust/src/scene.rs b/rust/src/scene.rs index 549476f..7b2e89d 100644 --- a/rust/src/scene.rs +++ b/rust/src/scene.rs @@ -423,8 +423,16 @@ impl SceneBuilder { pub fn build(self) -> Scene { let bounds: Vec = self.objects.iter().map(|o| o.shape.bounds().unwrap()).collect(); + let spheres: Vec> = self + .objects + .iter() + .map(|o| match o.shape { + Shape::Sphere { center, radius } => Some((center, radius)), + _ => None, + }) + .collect(); Scene { - bvh: Bvh::build(&bounds), + bvh: Bvh::build(&bounds, &spheres), objects: self.objects, planes: self.planes, lights: self.lights, @@ -448,6 +456,14 @@ impl Scene { self.objects.len() + self.planes.len() } + pub fn bvh_depth(&self) -> usize { + self.bvh.depth() + } + + pub fn uses_avx512(&self) -> bool { + self.bvh.uses_avx512() + } + pub fn light_count(&self) -> usize { self.lights.len() } @@ -488,12 +504,22 @@ impl Scene { closest = Some((t, outward, obj.material)); } } - self.bvh.closest(r, t_min, min_t, |i, t_max| { - let obj = &self.objects[i]; - let (t, outward) = obj.shape.hit(r, t_min, t_max)?; - closest = Some((t, outward, obj.material)); + // Spheres are intersected inside the BVH (8 at a time); other shapes + // come back through this callback, which remembers their normal. + let mut other_normal = Vec3::ZERO; + let bvh_hit = self.bvh.closest(r, t_min, min_t, |i, t_max| { + let (t, outward) = self.objects[i].shape.hit(r, t_min, t_max)?; + other_normal = outward; Some(t) }); + if let Some((i, t, is_sphere)) = bvh_hit { + let obj = &self.objects[i]; + let outward = match obj.shape { + Shape::Sphere { center, .. } if is_sphere => (r.at(t) - center).normalize(), + _ => other_normal, + }; + closest = Some((t, outward, obj.material)); + } closest.map(|(t, outward, material)| HitRecord::new(r, t, outward, material)) }