broodjeaap/ai_pathtracer
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

simd AVX512

Browse source
This commit is contained in:
David 2026-09-26 11:31:08 +02:00
parent e5d3c6e4e1
commit a86a2a64f7
4 changed files with 665 additions and 132 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
README.md
View file

@ -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 | | `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 | | `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 | | `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. 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 ### 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. - **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. - **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.

752
rust/src/bvh.rs
View file

@ -1,5 +1,15 @@
//! Bounding volume hierarchy over the scene's bounded objects, so a ray only //! 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. //! 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::scene::Ray;
use crate::vec3::Vec3; use crate::vec3::Vec3;
@ -27,159 +37,651 @@ impl Aabb {
(self.min + self.max) * 0.5 (self.min + self.max) * 0.5
} }
/// Slab test. `inv_dir` components may be infinite for axis-parallel rays; fn surface_area(&self) -> f64 {
/// the resulting NaNs are discarded by f64::min/max. let d = self.max - self.min;
fn hit(&self, origin: Vec3, inv_dir: Vec3, mut t_min: f64, mut t_max: f64) -> bool { 2.0 * (d.x * d.y + d.y * d.z + d.x * d.z)
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
} }
} }
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`. // Build: binary BVH first.
struct Node {
struct BinNode {
bounds: Aabb, bounds: Aabb,
start: u32, start: usize,
count: u32, count: usize, // > 0 for leaves
right: u32, left: usize,
axis: u8, 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 { pub struct Bvh {
nodes: Vec<Node>, /// Reference to the root: node 0, or a single leaf when the whole scene
/// Object indices, reordered so each leaf's objects are contiguous. /// fits in one (then there is no node test at all). None if empty.
prims: Vec<u32>, 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 { impl Bvh {
pub fn build(bounds: &[Aabb]) -> Bvh { /// `spheres[i]` is Some((center, radius)) if object i is a sphere.
let mut bvh = Bvh { pub fn build(bounds: &[Aabb], spheres: &[Option<(Vec3, f64)>]) -> Bvh {
nodes: Vec::with_capacity(2 * bounds.len()), let mut slots: Vec<u32> = (0..bounds.len() as u32).collect();
prims: (0..bounds.len() as u32).collect(), let mut bin = Vec::with_capacity(2 * bounds.len());
};
if !bounds.is_empty() { if !bounds.is_empty() {
let centroids: Vec<Vec3> = bounds.iter().map(Aabb::centroid).collect(); let input = BuildInput {
bvh.build_range(0, bounds.len(), bounds, &centroids); 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 bvh
} }
/// Median split along the axis where the object centroids spread widest. fn leaf(&self, b: &BinNode) -> u32 {
fn build_range(&mut self, start: usize, end: usize, bounds: &[Aabb], centroids: &[Vec3]) -> usize { let r2 = &self.spheres.r2[b.start..b.start + b.count];
let prims = &self.prims[start..end]; let has_spheres = r2.iter().any(|&r| r != f64::NEG_INFINITY);
let node_bounds = prims.iter().fold(Aabb::EMPTY, |b, &p| b.union(bounds[p as usize])); let has_others = r2.iter().any(|&r| r == f64::NEG_INFINITY);
let centroid_bounds = prims.iter().fold(Aabb::EMPTY, |b, &p| { leaf_ref(b.start, b.count, has_spheres, has_others)
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 idx = self.nodes.len(); /// Converts binary node `b` (and its subtree) into 8-wide nodes, returning
self.nodes.push(Node { /// the new node's index. Starts from b's two children and repeatedly
bounds: node_bounds, /// opens up the interior child with the largest surface area until there
start: start as u32, /// are 8 children or only leaves remain.
count: (end - start) as u32, fn collapse(&mut self, bin: &[BinNode], b: usize) -> usize {
right: 0, let mut children = vec![bin[b].left, bin[b].right];
axis: axis as u8, while children.len() < WIDTH {
}); let Some(pos) = children
if end - start <= LEAF_SIZE || extent[axis] <= 0.0 { .iter()
return idx; .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; let idx = self.nodes.len();
self.prims[start..end].select_nth_unstable_by(half, |&a, &b| { self.nodes.push(Node8::empty());
centroids[a as usize][axis].total_cmp(&centroids[b as usize][axis]) for (lane, &c) in children.iter().enumerate() {
}); let child_ref = if bin[c].count > 0 {
self.build_range(start, start + half, bounds, centroids); self.leaf(&bin[c])
let right = self.build_range(start + half, end, bounds, centroids); } else {
self.nodes[idx].count = 0; self.collapse(bin, c) as u32
self.nodes[idx].right = right 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 idx
} }
/// Calls `visit` for each object in a leaf whose box the ray passes through /// Number of 8-wide node levels (0 when the root is a single leaf).
/// within (t_min, t_max). `visit` returns Some(t) to shrink t_max after a pub fn depth(&self) -> usize {
/// hit, or Err(()) to stop traversal entirely. fn go(b: &Bvh, n: usize) -> usize {
fn traverse( 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, &self,
r: &Ray, r: &Ray,
t_min: f64, t_min: f64,
mut t_max: f64, t_max: f64,
mut visit: impl FnMut(usize, f64) -> Result<Option<f64>, ()>, other: impl FnMut(usize, f64) -> Option<f64>,
) { ) -> Option<(usize, f64, bool)> {
if self.nodes.is_empty() { let Some(root) = self.root else {
return; 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); #[cfg(target_arch = "x86_64")]
let mut stack = [0u32; 64]; if self.simd_pays_off(root) {
let mut sp = 1; // SAFETY: simd_pays_off() requires use_avx512, which is only set when the CPU supports AVX-512F.
while sp > 0 { return unsafe { self.closest_avx512(root, &rd, t_min, t_max, other) };
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;
}
} }
unsafe { self.closest_scalar(root, &rd, t_min, t_max, other) }
} }
/// Closest hit: `hit(object, t_max)` returns the hit distance if the object pub fn any(&self, r: &Ray, t_min: f64, t_max: f64, other: impl FnMut(usize) -> bool) -> bool {
/// is hit closer than t_max. let Some(root) = self.root else {
pub fn closest(&self, r: &Ray, t_min: f64, t_max: f64, mut hit: impl FnMut(usize, f64) -> Option<f64>) { return false;
self.traverse(r, t_min, t_max, |p, t_max| Ok(hit(p, t_max))); };
} let rd = RayData::new(r);
let t_max = t_max.min(f64::MAX); // See Node8.
/// True as soon as `hit(object)` reports any hit. if !self.root_hit(root, &rd, t_min, t_max) {
pub fn any(&self, r: &Ray, t_min: f64, t_max: f64, mut hit: impl FnMut(usize) -> bool) -> bool { return false;
let mut found = false; }
self.traverse(r, t_min, t_max, |p, _| { #[cfg(target_arch = "x86_64")]
if hit(p) { if self.simd_pays_off(root) {
found = true; // SAFETY: as in closest().
Err(()) return unsafe { self.any_avx512(root, &rd, t_min, t_max, other) };
} else { }
Ok(None) unsafe { self.any_scalar(root, &rd, t_min, t_max, other) }
} }
}); }
found
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
} }
} }

5
rust/src/main.rs
View file

@ -409,6 +409,11 @@ fn main() {
println!("Adaptive sampling: off"); println!("Adaptive sampling: off");
} }
println!("Threads: {}", rayon::current_num_threads()); 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!("Output: {}", settings.output);
println!(); println!();

36
rust/src/scene.rs
View file

@ -423,8 +423,16 @@ impl SceneBuilder {
pub fn build(self) -> Scene { pub fn build(self) -> Scene {
let bounds: Vec<Aabb> = self.objects.iter().map(|o| o.shape.bounds().unwrap()).collect(); 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 { Scene {
bvh: Bvh::build(&bounds), bvh: Bvh::build(&bounds, &spheres),
objects: self.objects, objects: self.objects,
planes: self.planes, planes: self.planes,
lights: self.lights, lights: self.lights,
@ -448,6 +456,14 @@ impl Scene {
self.objects.len() + self.planes.len() 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 { pub fn light_count(&self) -> usize {
self.lights.len() self.lights.len()
} }
@ -488,12 +504,22 @@ impl Scene {
closest = Some((t, outward, obj.material)); closest = Some((t, outward, obj.material));
} }
} }
self.bvh.closest(r, t_min, min_t, |i, t_max| { // Spheres are intersected inside the BVH (8 at a time); other shapes
let obj = &self.objects[i]; // come back through this callback, which remembers their normal.
let (t, outward) = obj.shape.hit(r, t_min, t_max)?; let mut other_normal = Vec3::ZERO;
closest = Some((t, outward, obj.material)); 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) 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)) closest.map(|(t, outward, material)| HitRecord::new(r, t, outward, material))
} }