simd AVX512

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.

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<Vec3>,
+    is_sphere: Vec<bool>,
+}
+
+fn build_binary(nodes: &mut Vec<BinNode>, 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(&centroids[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<f64>,
+    cy: Vec<f64>,
+    cz: Vec<f64>,
+    r2: Vec<f64>,
+}
+
+/// 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<f64>| _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<f64>,
+        ) -> 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::<u32>::uninit(); STACK_SIZE];
+            let mut stack_t = [MaybeUninit::<f64>::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::<u32>::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<Node>,
-    /// Object indices, reordered so each leaf's objects are contiguous.
-    prims: Vec<u32>,
+    /// 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<u32>,
+    /// 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<Node8>,
+    /// Object index for each slot; each leaf covers a contiguous slot range.
+    slots: Vec<u32>,
+    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<u32> = (0..bounds.len() as u32).collect();
+        let mut bin = Vec::with_capacity(2 * bounds.len());
         if !bounds.is_empty() {
-            let centroids: Vec<Vec3> = bounds.iter().map(Aabb::centroid).collect();
-            bvh.build_range(0, bounds.len(), bounds, &centroids);
+            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(&centroids[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<Option<f64>, ()>,
-    ) {
-        if self.nodes.is_empty() {
-            return;
+        t_max: f64,
+        other: impl FnMut(usize, f64) -> Option<f64>,
+    ) -> 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<f64>) {
-        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
     }
 }

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!();
 

rust/src/scene.rs

@@ -423,8 +423,16 @@ impl SceneBuilder {
 
     pub fn build(self) -> Scene {
         let bounds: Vec<Aabb> = self.objects.iter().map(|o| o.shape.bounds().unwrap()).collect();
+        let spheres: Vec<Option<(Vec3, f64)>> = 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))
     }